1use alloc::borrow::Cow;
4use alloc::fmt;
5use alloc::vec::Vec;
6use std::error;
7use std::io;
8use std::io::Write;
9
10use weezl::{encode::Encoder as LzwEncoder, BitOrder};
11
12use crate::common::{AnyExtension, Block, DisposalMethod, Extension, Frame};
13use crate::traits::WriteBytesExt;
14
15#[derive(Debug)]
17#[non_exhaustive]
18pub enum EncodingFormatError {
19 TooManyColors,
21 MissingColorPalette,
23 InvalidMinCodeSize,
25}
26
27impl error::Error for EncodingFormatError {}
28impl fmt::Display for EncodingFormatError {
29 #[cold]
30 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
31 match self {
32 Self::TooManyColors => write!(fmt, "the image has too many colors"),
33 Self::MissingColorPalette => write!(
34 fmt,
35 "the GIF format requires a color palette but none was given"
36 ),
37 Self::InvalidMinCodeSize => write!(fmt, "LZW data is invalid"),
38 }
39 }
40}
41
42#[derive(Debug)]
44#[non_exhaustive]
45pub enum EncodingError {
46 FrameBufferTooSmallForDimensions,
48 OutOfMemory,
50 WriterNotFound,
52 Format(EncodingFormatError),
54 Io(io::Error),
56}
57
58impl fmt::Display for EncodingError {
59 #[cold]
60 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
61 match self {
62 Self::FrameBufferTooSmallForDimensions => {
63 fmt.write_str("Frame Buffer Too Small for Dimensions")
64 }
65 Self::OutOfMemory => fmt.write_str("Out of Memory"),
66 Self::WriterNotFound => fmt.write_str("Writer Not Found"),
67 Self::Io(err) => err.fmt(fmt),
68 Self::Format(err) => err.fmt(fmt),
69 }
70 }
71}
72
73impl error::Error for EncodingError {
74 #[cold]
75 fn source(&self) -> Option<&(dyn error::Error + 'static)> {
76 match self {
77 Self::FrameBufferTooSmallForDimensions => None,
78 Self::OutOfMemory => None,
79 Self::WriterNotFound => None,
80 Self::Io(err) => Some(err),
81 Self::Format(err) => Some(err),
82 }
83 }
84}
85
86impl From<io::Error> for EncodingError {
87 #[cold]
88 fn from(err: io::Error) -> Self {
89 Self::Io(err)
90 }
91}
92
93impl From<EncodingFormatError> for EncodingError {
94 #[cold]
95 fn from(err: EncodingFormatError) -> Self {
96 Self::Format(err)
97 }
98}
99
100#[derive(Copy, Clone, Debug, PartialEq, Eq)]
102pub enum Repeat {
103 Finite(u16),
105 Infinite,
107}
108
109impl Default for Repeat {
110 fn default() -> Self {
111 Self::Finite(0)
112 }
113}
114
115#[non_exhaustive]
117pub enum ExtensionData {
118 Control {
120 flags: u8,
122 delay: u16,
124 trns: u8,
126 },
127 Repetitions(Repeat),
129}
130
131impl ExtensionData {
132 #[must_use]
136 pub fn new_control_ext(
137 delay: u16,
138 dispose: DisposalMethod,
139 needs_user_input: bool,
140 trns: Option<u8>,
141 ) -> Self {
142 let mut flags = 0;
143 let trns = match trns {
144 Some(trns) => {
145 flags |= 1;
146 trns
147 }
148 None => 0,
149 };
150 flags |= u8::from(needs_user_input) << 1;
151 flags |= (dispose as u8) << 2;
152 Self::Control { flags, delay, trns }
153 }
154}
155
156impl<W: Write> Encoder<W> {
157 pub fn new(
162 w: W,
163 width: u16,
164 height: u16,
165 global_palette: &[u8],
166 ) -> Result<Self, EncodingError> {
167 Self {
168 w: Some(w),
169 global_palette: false,
170 width,
171 height,
172 buffer: Vec::new(),
173 }
174 .write_global_palette(global_palette)
175 }
176
177 pub fn set_repeat(&mut self, repeat: Repeat) -> Result<(), EncodingError> {
179 self.write_extension(ExtensionData::Repetitions(repeat))
180 }
181
182 fn write_global_palette(mut self, palette: &[u8]) -> Result<Self, EncodingError> {
184 let mut flags = 0;
185 flags |= 0b1000_0000;
186 let (palette, padding, table_size) = Self::check_color_table(palette)?;
187 self.global_palette = !palette.is_empty();
188 flags |= table_size;
190 flags |= table_size << 4; self.write_screen_desc(flags)?;
194 Self::write_color_table(self.writer()?, palette, padding)?;
195 Ok(self)
196 }
197
198 pub fn write_frame(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
202 if usize::from(frame.width)
203 .checked_mul(usize::from(frame.height))
204 .map_or(true, |size| frame.buffer.len() < size)
205 {
206 return Err(EncodingError::FrameBufferTooSmallForDimensions);
207 }
208 debug_assert!(
209 (frame.width > 0 && frame.height > 0) || frame.buffer.is_empty(),
210 "the frame has 0 pixels, but non-empty buffer"
211 );
212 self.write_frame_header(frame)?;
213 self.write_image_block(&frame.buffer)
214 }
215
216 fn write_frame_header(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
217 self.write_extension(ExtensionData::new_control_ext(
218 frame.delay,
219 frame.dispose,
220 frame.needs_user_input,
221 frame.transparent,
222 ))?;
223 let mut flags = 0;
224 if frame.interlaced {
225 flags |= 0b0100_0000;
226 }
227 let palette = match frame.palette {
228 Some(ref palette) => {
229 flags |= 0b1000_0000;
230 let (palette, padding, table_size) = Self::check_color_table(palette)?;
231 flags |= table_size;
232 Some((palette, padding))
233 }
234 None if self.global_palette => None,
235 _ => {
236 return Err(EncodingError::from(
237 EncodingFormatError::MissingColorPalette,
238 ))
239 }
240 };
241 let mut tmp = tmp_buf::<10>();
242 tmp.write_le(Block::Image as u8)?;
243 tmp.write_le(frame.left)?;
244 tmp.write_le(frame.top)?;
245 tmp.write_le(frame.width)?;
246 tmp.write_le(frame.height)?;
247 tmp.write_le(flags)?;
248 let writer = self.writer()?;
249 tmp.finish(&mut *writer)?;
250 if let Some((palette, padding)) = palette {
251 Self::write_color_table(writer, palette, padding)?;
252 }
253 Ok(())
254 }
255
256 fn write_image_block(&mut self, data: &[u8]) -> Result<(), EncodingError> {
257 self.buffer.clear();
258 self.buffer
259 .try_reserve(data.len() / 4)
260 .map_err(|_| EncodingError::OutOfMemory)?;
261 lzw_encode(data, &mut self.buffer);
262
263 let writer = self.w.as_mut().ok_or(EncodingError::WriterNotFound)?;
264 Self::write_encoded_image_block(writer, &self.buffer)
265 }
266
267 fn write_encoded_image_block(
268 writer: &mut W,
269 data_with_min_code_size: &[u8],
270 ) -> Result<(), EncodingError> {
271 let (&min_code_size, data) = data_with_min_code_size.split_first().unwrap_or((&2, &[]));
272 writer.write_le(min_code_size)?;
273
274 let mut iter = data.chunks_exact(0xFF);
277 for full_block in iter.by_ref() {
278 writer.write_le(0xFFu8)?;
279 writer.write_all(full_block)?;
280 }
281 let last_block = iter.remainder();
282 if !last_block.is_empty() {
283 writer.write_le(last_block.len() as u8)?;
284 writer.write_all(last_block)?;
285 }
286 writer.write_le(0u8).map_err(Into::into)
287 }
288
289 fn write_color_table(
290 writer: &mut W,
291 table: &[u8],
292 padding: usize,
293 ) -> Result<(), EncodingError> {
294 writer.write_all(table)?;
295 for _ in 0..padding {
297 writer.write_all(&[0, 0, 0])?;
298 }
299 Ok(())
300 }
301
302 fn check_color_table(table: &[u8]) -> Result<(&[u8], usize, u8), EncodingError> {
304 let num_colors = table.len() / 3;
305 if num_colors > 256 {
306 return Err(EncodingError::from(EncodingFormatError::TooManyColors));
307 }
308 let table_size = flag_size(num_colors);
309 let padding = (2 << table_size) - num_colors;
310 Ok((&table[..num_colors * 3], padding, table_size))
311 }
312
313 pub fn write_extension(&mut self, extension: ExtensionData) -> Result<(), EncodingError> {
317 use self::ExtensionData::*;
318 if let Repetitions(Repeat::Finite(0)) = extension {
321 return Ok(());
322 }
323 let writer = self.writer()?;
324 writer.write_le(Block::Extension as u8)?;
325 match extension {
326 Control { flags, delay, trns } => {
327 let mut tmp = tmp_buf::<6>();
328 tmp.write_le(Extension::Control as u8)?;
329 tmp.write_le(4u8)?;
330 tmp.write_le(flags)?;
331 tmp.write_le(delay)?;
332 tmp.write_le(trns)?;
333 tmp.finish(&mut *writer)?;
334 }
335 Repetitions(repeat) => {
336 let mut tmp = tmp_buf::<17>();
337 tmp.write_le(Extension::Application as u8)?;
338 tmp.write_le(11u8)?;
339 tmp.write_all(b"NETSCAPE2.0")?;
340 tmp.write_le(3u8)?;
341 tmp.write_le(1u8)?;
342 tmp.write_le(match repeat {
343 Repeat::Finite(no) => no,
344 Repeat::Infinite => 0u16,
345 })?;
346 tmp.finish(&mut *writer)?;
347 }
348 }
349 writer.write_le(0u8).map_err(Into::into)
350 }
351
352 pub fn write_raw_extension(
358 &mut self,
359 func: AnyExtension,
360 data: &[&[u8]],
361 ) -> Result<(), EncodingError> {
362 let writer = self.writer()?;
363 writer.write_le(Block::Extension as u8)?;
364 writer.write_le(func.0)?;
365 for block in data {
366 for chunk in block.chunks(0xFF) {
367 writer.write_le(chunk.len() as u8)?;
368 writer.write_all(chunk)?;
369 }
370 }
371 Ok(writer.write_le(0u8)?)
372 }
373
374 pub fn write_lzw_pre_encoded_frame(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
379 if let Some(&min_code_size) = frame.buffer.first() {
381 if min_code_size > 11 || min_code_size < 2 {
382 return Err(EncodingError::Format(
383 EncodingFormatError::InvalidMinCodeSize,
384 ));
385 }
386 }
387
388 self.write_frame_header(frame)?;
389 let writer = self.writer()?;
390 Self::write_encoded_image_block(writer, &frame.buffer)
391 }
392
393 fn write_screen_desc(&mut self, flags: u8) -> Result<(), EncodingError> {
395 let mut tmp = tmp_buf::<13>();
396 tmp.write_all(b"GIF89a")?;
397 tmp.write_le(self.width)?;
398 tmp.write_le(self.height)?;
399 tmp.write_le(flags)?; tmp.write_le(0u8)?; tmp.write_le(0u8)?; Ok(tmp.finish(self.writer()?)?)
403 }
404
405 pub fn get_ref(&self) -> &W {
407 self.w.as_ref().unwrap()
408 }
409
410 pub fn get_mut(&mut self) -> &mut W {
414 self.w.as_mut().unwrap()
415 }
416
417 pub fn into_inner(mut self) -> Result<W, EncodingError> {
419 self.write_trailer()?;
420 self.w.take().ok_or(EncodingError::WriterNotFound)
421 }
422
423 fn write_trailer(&mut self) -> Result<(), EncodingError> {
425 Ok(self.writer()?.write_le(Block::Trailer as u8)?)
426 }
427
428 #[inline]
429 fn writer(&mut self) -> Result<&mut W, EncodingError> {
430 self.w.as_mut().ok_or(EncodingError::WriterNotFound)
431 }
432}
433
434fn lzw_encode(data: &[u8], buffer: &mut Vec<u8>) {
438 let mut max_byte = 0;
439 for &byte in data {
440 if byte > max_byte {
441 max_byte = byte;
442 if byte > 127 {
444 break;
445 }
446 }
447 }
448 let palette_min_len = u32::from(max_byte) + 1;
449 let min_code_size = palette_min_len.max(4).next_power_of_two().trailing_zeros() as u8;
451 buffer.push(min_code_size);
452 let mut enc = LzwEncoder::new(BitOrder::Lsb, min_code_size);
453 let len = enc.into_vec(buffer).encode_all(data).consumed_out;
454 buffer.truncate(len + 1);
455}
456
457impl Frame<'_> {
458 pub fn make_lzw_pre_encoded(&mut self) {
462 let mut buffer = Vec::new();
463 buffer.try_reserve(self.buffer.len() / 2).expect("OOM");
464 lzw_encode(&self.buffer, &mut buffer);
465 self.buffer = Cow::Owned(buffer);
466 }
467}
468
469pub struct Encoder<W: Write> {
471 w: Option<W>,
472 global_palette: bool,
473 width: u16,
474 height: u16,
475 buffer: Vec<u8>,
476}
477
478impl<W: Write> Drop for Encoder<W> {
479 #[cfg(feature = "raii_no_panic")]
480 fn drop(&mut self) {
481 if self.w.is_some() {
482 let _ = self.write_trailer();
483 }
484 }
485
486 #[cfg(not(feature = "raii_no_panic"))]
487 fn drop(&mut self) {
488 if self.w.is_some() {
489 self.write_trailer().unwrap();
490 }
491 }
492}
493
494fn flag_size(size: usize) -> u8 {
496 (size.clamp(2, 255).next_power_of_two().trailing_zeros() - 1) as u8
497}
498
499#[test]
500fn test_flag_size() {
501 #[rustfmt::skip]
502 fn expected(size: usize) -> u8 {
503 match size {
504 0 ..=2 => 0,
505 3 ..=4 => 1,
506 5 ..=8 => 2,
507 9 ..=16 => 3,
508 17 ..=32 => 4,
509 33 ..=64 => 5,
510 65 ..=128 => 6,
511 129..=256 => 7,
512 _ => 7
513 }
514 }
515
516 for i in 0..300 {
517 assert_eq!(flag_size(i), expected(i));
518 }
519 for i in 4..=255u8 {
520 let expected = match flag_size(1 + i as usize) + 1 {
521 1 => 2,
522 n => n,
523 };
524 let actual = (u32::from(i) + 1)
525 .max(4)
526 .next_power_of_two()
527 .trailing_zeros() as u8;
528 assert_eq!(actual, expected);
529 }
530}
531
532struct Buf<const N: usize> {
533 buf: [u8; N],
534 pos: usize,
535}
536
537impl<const N: usize> Write for Buf<N> {
538 #[inline(always)]
539 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
540 let len = buf.len();
541 let pos = self.pos;
542 self.buf
543 .get_mut(pos..pos + len)
544 .ok_or(io::ErrorKind::WriteZero)?
545 .copy_from_slice(buf);
546 self.pos += len;
547 Ok(len)
548 }
549
550 fn flush(&mut self) -> io::Result<()> {
551 Ok(())
552 }
553}
554
555fn tmp_buf<const N: usize>() -> Buf<N> {
556 Buf {
557 buf: [0; N],
558 pos: 0,
559 }
560}
561
562impl<const N: usize> Buf<N> {
563 #[inline(always)]
564 fn finish(&self, mut w: impl Write) -> io::Result<()> {
565 debug_assert_eq!(self.pos, N);
566 w.write_all(&self.buf)
567 }
568}
569
570#[test]
571fn error_cast() {
572 use alloc::boxed::Box;
573 let _: Box<dyn error::Error> =
574 EncodingError::from(EncodingFormatError::MissingColorPalette).into();
575}