Skip to main content

gif/
encoder.rs

1//! # Minimal gif encoder
2
3use 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/// The image has incorrect properties, making it impossible to encode as a gif.
16#[derive(Debug)]
17#[non_exhaustive]
18pub enum EncodingFormatError {
19    /// The image has too many colors.
20    TooManyColors,
21    /// The image has no color palette which is required.
22    MissingColorPalette,
23    /// LZW data is not valid for GIF. This may happen when wrong buffer is given to `write_lzw_pre_encoded_frame`
24    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/// Encoding error.
43#[derive(Debug)]
44#[non_exhaustive]
45pub enum EncodingError {
46    /// Frame buffer is too small for the declared dimensions.
47    FrameBufferTooSmallForDimensions,
48    /// Failed to internally allocate a buffer of sufficient size.
49    OutOfMemory,
50    /// Expected a writer but none found.
51    WriterNotFound,
52    /// Returned if the to image is not encodable as a gif.
53    Format(EncodingFormatError),
54    /// Wraps `std::io::Error`.
55    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/// Number of repetitions
101#[derive(Copy, Clone, Debug, PartialEq, Eq)]
102pub enum Repeat {
103    /// Finite number of repetitions
104    Finite(u16),
105    /// Infinite number of repetitions
106    Infinite,
107}
108
109impl Default for Repeat {
110    fn default() -> Self {
111        Self::Finite(0)
112    }
113}
114
115/// Extension data.
116#[non_exhaustive]
117pub enum ExtensionData {
118    /// Control extension. Use `ExtensionData::new_control_ext` to construct.
119    Control {
120        /// Flags.
121        flags: u8,
122        /// Frame delay.
123        delay: u16,
124        /// Transparent index.
125        trns: u8,
126    },
127    /// Sets the number of repetitions
128    Repetitions(Repeat),
129}
130
131impl ExtensionData {
132    /// Constructor for control extension data.
133    ///
134    /// `delay` is given in units of 10 ms.
135    #[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    /// Creates a new encoder.
158    ///
159    /// `global_palette` gives the global color palette in the format `[r, g, b, ...]`,
160    /// if no global palette shall be used an empty slice may be supplied.
161    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    /// Write an extension block that signals a repeat behaviour.
178    pub fn set_repeat(&mut self, repeat: Repeat) -> Result<(), EncodingError> {
179        self.write_extension(ExtensionData::Repetitions(repeat))
180    }
181
182    /// Writes the global color palette.
183    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        // Size of global color table.
189        flags |= table_size;
190        // Color resolution .. FIXME. This is mostly ignored (by ImageMagick at least) but hey, we
191        // should use some sensible value here or even allow configuring it?
192        flags |= table_size << 4; // wtf flag
193        self.write_screen_desc(flags)?;
194        Self::write_color_table(self.writer()?, palette, padding)?;
195        Ok(self)
196    }
197
198    /// Writes a frame to the image.
199    ///
200    /// Note: This function also writes a control extension if necessary.
201    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        // Write blocks. `chunks_exact` seems to be slightly faster
275        // than `chunks` according to both Rust docs and benchmark results.
276        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        // Waste some space as of gif spec
296        for _ in 0..padding {
297            writer.write_all(&[0, 0, 0])?;
298        }
299        Ok(())
300    }
301
302    /// returns rounded palette size, number of missing colors, and table size flag
303    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    /// Writes an extension to the image.
314    ///
315    /// It is normally not necessary to call this method manually.
316    pub fn write_extension(&mut self, extension: ExtensionData) -> Result<(), EncodingError> {
317        use self::ExtensionData::*;
318        // 0 finite repetitions can only be achieved
319        // if the corresponting extension is not written
320        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    /// Writes a raw extension to the image.
353    ///
354    /// This method can be used to write an unsupported extension to the file. `func` is the extension
355    /// identifier (e.g. `Extension::Application as u8`). `data` are the extension payload blocks. If any
356    /// contained slice has a lenght > 255 it is automatically divided into sub-blocks.
357    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    /// Writes a frame to the image, but expects `Frame.buffer` to contain LZW-encoded data
375    /// from [`Frame::make_lzw_pre_encoded`].
376    ///
377    /// Note: This function also writes a control extension if necessary.
378    pub fn write_lzw_pre_encoded_frame(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
379        // empty data is allowed
380        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    /// Writes the logical screen desriptor
394    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)?; // packed field
400        tmp.write_le(0u8)?; // bg index
401        tmp.write_le(0u8)?; // aspect ratio
402        Ok(tmp.finish(self.writer()?)?)
403    }
404
405    /// Gets a reference to the writer instance used by this encoder.
406    pub fn get_ref(&self) -> &W {
407        self.w.as_ref().unwrap()
408    }
409
410    /// Gets a mutable reference to the writer instance used by this encoder.
411    ///
412    /// It is inadvisable to directly write to the underlying writer.
413    pub fn get_mut(&mut self) -> &mut W {
414        self.w.as_mut().unwrap()
415    }
416
417    /// Finishes writing, and returns the `io::Write` instance used by this encoder
418    pub fn into_inner(mut self) -> Result<W, EncodingError> {
419        self.write_trailer()?;
420        self.w.take().ok_or(EncodingError::WriterNotFound)
421    }
422
423    /// Write the final tailer.
424    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
434/// Encodes the data into the provided buffer.
435///
436/// The first byte is the minimum code size, followed by LZW data.
437fn 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            // code size is the same after that
443            if byte > 127 {
444                break;
445            }
446        }
447    }
448    let palette_min_len = u32::from(max_byte) + 1;
449    // As per gif spec: The minimal code size has to be >= 2
450    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    /// Replace frame's buffer with a LZW-compressed one for use with [`Encoder::write_lzw_pre_encoded_frame`].
459    ///
460    /// Frames can be compressed in any order, separately from the `Encoder`, which can be used to compress frames in parallel.
461    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
469/// GIF encoder.
470pub 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
494// Color table size converted to flag bits
495fn 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}