gif/
encoder.rs

1//! # Minimal gif encoder
2use std::io;
3use std::io::prelude::*;
4use std::fmt;
5use std::error;
6use std::borrow::Cow;
7
8use weezl::{BitOrder, encode::Encoder as LzwEncoder};
9
10use crate::traits::WriteBytesExt;
11use crate::common::{AnyExtension, Block, DisposalMethod, Extension, Frame};
12
13/// The image has incorrect properties, making it impossible to encode as a gif.
14#[derive(Debug)]
15#[non_exhaustive]
16pub enum EncodingFormatError {
17    /// The image has too many colors.
18    TooManyColors,
19    /// The image has no color palette which is required.
20    MissingColorPalette,
21    /// LZW data is not valid for GIF. This may happen when wrong buffer is given to `write_lzw_pre_encoded_frame`
22    InvalidMinCodeSize,
23}
24
25impl error::Error for EncodingFormatError {}
26impl fmt::Display for EncodingFormatError {
27    #[cold]
28    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
29        match self {
30            Self::TooManyColors => write!(fmt, "the image has too many colors"),
31            Self::MissingColorPalette => write!(fmt, "the GIF format requires a color palette but none was given"),
32            Self::InvalidMinCodeSize => write!(fmt, "LZW data is invalid"),
33        }
34    }
35}
36
37#[derive(Debug)]
38/// Encoding error.
39pub enum EncodingError {
40    /// Returned if the to image is not encodable as a gif.
41    Format(EncodingFormatError),
42    /// Wraps `std::io::Error`.
43    Io(io::Error),
44}
45
46impl fmt::Display for EncodingError {
47    #[cold]
48    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
49        match self {
50            EncodingError::Io(err) => err.fmt(fmt),
51            EncodingError::Format(err) => err.fmt(fmt),
52        }
53    }
54}
55
56impl error::Error for EncodingError {
57    #[cold]
58    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
59        match self {
60            EncodingError::Io(err) => Some(err),
61            EncodingError::Format(err) => Some(err),
62        }
63    }
64}
65
66impl From<io::Error> for EncodingError {
67    #[cold]
68    fn from(err: io::Error) -> Self {
69        EncodingError::Io(err)
70    }
71}
72
73impl From<EncodingFormatError> for EncodingError {
74    #[cold]
75    fn from(err: EncodingFormatError) -> Self {
76        EncodingError::Format(err)
77    }
78}
79
80/// Number of repetitions
81#[derive(Copy, Clone, Debug, PartialEq, Eq)]
82pub enum Repeat {
83    /// Finite number of repetitions
84    Finite(u16),
85    /// Infinite number of repetitions
86    Infinite,
87}
88
89impl Default for Repeat {
90    fn default() -> Self {
91        Self::Finite(0)
92    }
93}
94
95/// Extension data.
96#[non_exhaustive]
97pub enum ExtensionData {
98    /// Control extension. Use `ExtensionData::new_control_ext` to construct.
99    Control {
100        /// Flags.
101        flags: u8,
102        /// Frame delay.
103        delay: u16,
104        /// Transparent index.
105        trns: u8,
106    },
107    /// Sets the number of repetitions
108    Repetitions(Repeat),
109}
110
111impl ExtensionData {
112    /// Constructor for control extension data.
113    ///
114    /// `delay` is given in units of 10 ms.
115    #[must_use] pub fn new_control_ext(delay: u16, dispose: DisposalMethod,
116                           needs_user_input: bool, trns: Option<u8>) -> ExtensionData {
117        let mut flags = 0;
118        let trns = match trns {
119            Some(trns) => {
120                flags |= 1;
121                trns
122            },
123            None => 0
124        };
125        flags |= u8::from(needs_user_input) << 1;
126        flags |= (dispose as u8) << 2;
127        ExtensionData::Control { flags, delay, trns }
128    }
129}
130
131impl<W: Write> Encoder<W> {
132    /// Creates a new encoder.
133    ///
134    /// `global_palette` gives the global color palette in the format `[r, g, b, ...]`,
135    /// if no global palette shall be used an empty slice may be supplied.
136    pub fn new(w: W, width: u16, height: u16, global_palette: &[u8]) -> Result<Self, EncodingError> {
137        Encoder {
138            w: Some(w),
139            global_palette: false,
140            width, height,
141            buffer: Vec::new(),
142        }.write_global_palette(global_palette)
143    }
144
145    /// Write an extension block that signals a repeat behaviour.
146    pub fn set_repeat(&mut self, repeat: Repeat) -> Result<(), EncodingError> {
147        self.write_extension(ExtensionData::Repetitions(repeat))
148    }
149
150    /// Writes the global color palette.
151    fn write_global_palette(mut self, palette: &[u8]) -> Result<Self, EncodingError> {
152        let mut flags = 0;
153        flags |= 0b1000_0000;
154        let (palette, padding, table_size) = Self::check_color_table(palette)?;
155        self.global_palette = !palette.is_empty();
156        // Size of global color table.
157        flags |= table_size;
158        // Color resolution .. FIXME. This is mostly ignored (by ImageMagick at least) but hey, we
159        // should use some sensible value here or even allow configuring it?
160        flags |= table_size << 4; // wtf flag
161        self.write_screen_desc(flags)?;
162        Self::write_color_table(self.writer()?, palette, padding)?;
163        Ok(self)
164    }
165
166    /// Writes a frame to the image.
167    ///
168    /// Note: This function also writes a control extension if necessary.
169    pub fn write_frame(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
170        if usize::from(frame.width).checked_mul(usize::from(frame.height)).map_or(true, |size| frame.buffer.len() < size) {
171            return Err(io::Error::new(io::ErrorKind::InvalidInput, "frame.buffer is too small for its width/height").into());
172        }
173        debug_assert!((frame.width > 0 && frame.height > 0) || frame.buffer.is_empty(), "the frame has 0 pixels, but non-empty buffer");
174        self.write_frame_header(frame)?;
175        self.write_image_block(&frame.buffer)
176    }
177
178    fn write_frame_header(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
179        self.write_extension(ExtensionData::new_control_ext(
180            frame.delay,
181            frame.dispose,
182            frame.needs_user_input,
183            frame.transparent,
184        ))?;
185        let mut flags = 0;
186        if frame.interlaced {
187            flags |= 0b0100_0000;
188        }
189        let palette = match frame.palette {
190            Some(ref palette) => {
191                flags |= 0b1000_0000;
192                let (palette, padding, table_size) = Self::check_color_table(&palette)?;
193                flags |= table_size;
194                Some((palette, padding))
195            },
196            None if self.global_palette => None,
197            _ => return Err(EncodingError::from(EncodingFormatError::MissingColorPalette))
198        };
199        let mut tmp = tmp_buf::<10>();
200        tmp.write_le(Block::Image as u8)?;
201        tmp.write_le(frame.left)?;
202        tmp.write_le(frame.top)?;
203        tmp.write_le(frame.width)?;
204        tmp.write_le(frame.height)?;
205        tmp.write_le(flags)?;
206        let writer = self.writer()?;
207        tmp.finish(&mut *writer)?;
208        if let Some((palette, padding)) = palette {
209            Self::write_color_table(writer, palette, padding)?;
210        }
211        Ok(())
212    }
213
214    fn write_image_block(&mut self, data: &[u8]) -> Result<(), EncodingError> {
215        self.buffer.clear();
216        self.buffer.try_reserve(data.len() / 4)
217            .map_err(|_| io::Error::from(io::ErrorKind::OutOfMemory))?;
218        lzw_encode(data, &mut self.buffer);
219
220        let writer = self.w.as_mut().ok_or(io::Error::from(io::ErrorKind::Unsupported))?;
221        Self::write_encoded_image_block(writer, &self.buffer)
222    }
223
224    fn write_encoded_image_block(writer: &mut W, data_with_min_code_size: &[u8]) -> Result<(), EncodingError> {
225        let (&min_code_size, data) = data_with_min_code_size.split_first().unwrap_or((&2, &[]));
226        writer.write_le(min_code_size)?;
227
228        // Write blocks. `chunks_exact` seems to be slightly faster
229        // than `chunks` according to both Rust docs and benchmark results.
230        let mut iter = data.chunks_exact(0xFF);
231        for full_block in iter.by_ref() {
232            writer.write_le(0xFFu8)?;
233            writer.write_all(full_block)?;
234        }
235        let last_block = iter.remainder();
236        if !last_block.is_empty() {
237            writer.write_le(last_block.len() as u8)?;
238            writer.write_all(last_block)?;
239        }
240        writer.write_le(0u8).map_err(Into::into)
241    }
242
243    fn write_color_table(writer: &mut W, table: &[u8], padding: usize) -> Result<(), EncodingError> {
244        writer.write_all(&table)?;
245        // Waste some space as of gif spec
246        for _ in 0..padding {
247            writer.write_all(&[0, 0, 0])?;
248        }
249        Ok(())
250    }
251
252    /// returns rounded palette size, number of missing colors, and table size flag
253    fn check_color_table(table: &[u8]) -> Result<(&[u8], usize, u8), EncodingError> {
254        let num_colors = table.len() / 3;
255        if num_colors > 256 {
256            return Err(EncodingError::from(EncodingFormatError::TooManyColors));
257        }
258        let table_size = flag_size(num_colors);
259        let padding = (2 << table_size) - num_colors;
260        Ok((&table[..num_colors * 3], padding, table_size))
261    }
262
263    /// Writes an extension to the image.
264    ///
265    /// It is normally not necessary to call this method manually.
266    pub fn write_extension(&mut self, extension: ExtensionData) -> Result<(), EncodingError> {
267        use self::ExtensionData::*;
268        // 0 finite repetitions can only be achieved
269        // if the corresponting extension is not written
270        if let Repetitions(Repeat::Finite(0)) = extension {
271            return Ok(());
272        }
273        let writer = self.writer()?;
274        writer.write_le(Block::Extension as u8)?;
275        match extension {
276            Control { flags, delay, trns } => {
277                let mut tmp = tmp_buf::<6>();
278                tmp.write_le(Extension::Control as u8)?;
279                tmp.write_le(4u8)?;
280                tmp.write_le(flags)?;
281                tmp.write_le(delay)?;
282                tmp.write_le(trns)?;
283                tmp.finish(&mut *writer)?;
284            }
285            Repetitions(repeat) => {
286                let mut tmp = tmp_buf::<17>();
287                tmp.write_le(Extension::Application as u8)?;
288                tmp.write_le(11u8)?;
289                tmp.write_all(b"NETSCAPE2.0")?;
290                tmp.write_le(3u8)?;
291                tmp.write_le(1u8)?;
292                tmp.write_le(match repeat {
293                    Repeat::Finite(no) => no,
294                    Repeat::Infinite => 0u16,
295                })?;
296                tmp.finish(&mut *writer)?;
297            }
298        }
299        writer.write_le(0u8).map_err(Into::into)
300    }
301
302    /// Writes a raw extension to the image.
303    ///
304    /// This method can be used to write an unsupported extension to the file. `func` is the extension
305    /// identifier (e.g. `Extension::Application as u8`). `data` are the extension payload blocks. If any
306    /// contained slice has a lenght > 255 it is automatically divided into sub-blocks.
307    pub fn write_raw_extension(&mut self, func: AnyExtension, data: &[&[u8]]) -> io::Result<()> {
308        let writer = self.writer()?;
309        writer.write_le(Block::Extension as u8)?;
310        writer.write_le(func.0)?;
311        for block in data {
312            for chunk in block.chunks(0xFF) {
313                writer.write_le(chunk.len() as u8)?;
314                writer.write_all(chunk)?;
315            }
316        }
317        writer.write_le(0u8)
318    }
319
320    /// Writes a frame to the image, but expects `Frame.buffer` to contain LZW-encoded data
321    /// from [`Frame::make_lzw_pre_encoded`].
322    ///
323    /// Note: This function also writes a control extension if necessary.
324    pub fn write_lzw_pre_encoded_frame(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
325        // empty data is allowed
326        if let Some(&min_code_size) = frame.buffer.get(0) {
327            if min_code_size > 11 || min_code_size < 2 {
328                return Err(EncodingError::Format(EncodingFormatError::InvalidMinCodeSize));
329            }
330        }
331
332        self.write_frame_header(frame)?;
333        let writer = self.writer()?;
334        Self::write_encoded_image_block(writer, &frame.buffer)
335    }
336
337    /// Writes the logical screen desriptor
338    fn write_screen_desc(&mut self, flags: u8) -> io::Result<()> {
339        let mut tmp = tmp_buf::<13>();
340        tmp.write_all(b"GIF89a")?;
341        tmp.write_le(self.width)?;
342        tmp.write_le(self.height)?;
343        tmp.write_le(flags)?; // packed field
344        tmp.write_le(0u8)?; // bg index
345        tmp.write_le(0u8)?; // aspect ratio
346        tmp.finish(self.writer()?)
347    }
348
349    /// Gets a reference to the writer instance used by this encoder.
350    pub fn get_ref(&self) -> &W {
351        self.w.as_ref().unwrap()
352    }
353
354    /// Gets a mutable reference to the writer instance used by this encoder.
355    ///
356    /// It is inadvisable to directly write to the underlying writer.
357    pub fn get_mut(&mut self) -> &mut W {
358        self.w.as_mut().unwrap()
359    }
360
361    /// Finishes writing, and returns the `io::Write` instance used by this encoder
362    pub fn into_inner(mut self) -> io::Result<W> {
363        self.write_trailer()?;
364        self.w.take().ok_or(io::Error::from(io::ErrorKind::Unsupported))
365    }
366
367    /// Write the final tailer.
368    fn write_trailer(&mut self) -> io::Result<()> {
369        self.writer()?.write_le(Block::Trailer as u8)
370    }
371
372    #[inline]
373    fn writer(&mut self) -> io::Result<&mut W> {
374        self.w.as_mut().ok_or(io::Error::from(io::ErrorKind::Unsupported))
375    }
376}
377
378/// Encodes the data into the provided buffer.
379///
380/// The first byte is the minimum code size, followed by LZW data.
381fn lzw_encode(data: &[u8], buffer: &mut Vec<u8>) {
382    let mut max_byte = 0;
383    for &byte in data {
384        if byte > max_byte {
385            max_byte = byte;
386            // code size is the same after that
387            if byte > 127 {
388                break;
389            }
390        }
391    }
392    let palette_min_len = max_byte as u32 + 1;
393    // As per gif spec: The minimal code size has to be >= 2
394    let min_code_size = palette_min_len.max(4).next_power_of_two().trailing_zeros() as u8;
395    buffer.push(min_code_size);
396    let mut enc = LzwEncoder::new(BitOrder::Lsb, min_code_size);
397    let len = enc.into_vec(buffer).encode_all(data).consumed_out;
398    buffer.truncate(len+1);
399}
400
401impl Frame<'_> {
402    /// Replace frame's buffer with a LZW-compressed one for use with [`Encoder::write_lzw_pre_encoded_frame`].
403    ///
404    /// Frames can be compressed in any order, separately from the `Encoder`, which can be used to compress frames in parallel.
405    pub fn make_lzw_pre_encoded(&mut self) {
406        let mut buffer = Vec::new();
407        buffer.try_reserve(self.buffer.len() / 2).expect("OOM");
408        lzw_encode(&self.buffer, &mut buffer);
409        self.buffer = Cow::Owned(buffer);
410    }
411}
412
413/// GIF encoder.
414pub struct Encoder<W: Write> {
415    w: Option<W>,
416    global_palette: bool,
417    width: u16,
418    height: u16,
419    buffer: Vec<u8>,
420}
421
422impl<W: Write> Drop for Encoder<W> {
423    #[cfg(feature = "raii_no_panic")]
424    fn drop(&mut self) {
425        if self.w.is_some() {
426            let _ = self.write_trailer();
427        }
428    }
429
430    #[cfg(not(feature = "raii_no_panic"))]
431    fn drop(&mut self) {
432        if self.w.is_some() {
433            self.write_trailer().unwrap();
434        }
435    }
436}
437
438// Color table size converted to flag bits
439fn flag_size(size: usize) -> u8 {
440    (size.max(2).min(255).next_power_of_two().trailing_zeros()-1) as u8
441}
442
443#[test]
444fn test_flag_size() {
445    fn expected(size: usize) -> u8 {
446        match size {
447            0  ..=2   => 0,
448            3  ..=4   => 1,
449            5  ..=8   => 2,
450            9  ..=16  => 3,
451            17 ..=32  => 4,
452            33 ..=64  => 5,
453            65 ..=128 => 6,
454            129..=256 => 7,
455            _ => 7
456        }
457    }
458
459    for i in 0..300 {
460        assert_eq!(flag_size(i), expected(i));
461    }
462    for i in 4..=255u8 {
463        let expected = match flag_size(1 + i as usize) + 1 {
464            1 => 2,
465            n => n,
466        };
467        let actual = (i as u32 + 1).max(4).next_power_of_two().trailing_zeros() as u8;
468        assert_eq!(actual, expected);
469    }
470}
471
472struct Buf<const N: usize> {
473    buf: [u8; N], pos: usize
474}
475
476impl<const N: usize> Write for Buf<N> {
477    #[inline(always)]
478    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
479        let len = buf.len();
480        let pos = self.pos;
481        self.buf[pos.. pos + len].copy_from_slice(buf);
482        self.pos += len;
483        Ok(len)
484    }
485
486    fn flush(&mut self) -> io::Result<()> { Ok(()) }
487}
488
489fn tmp_buf<const N: usize>() -> Buf<N> {
490    Buf { buf: [0; N], pos: 0 }
491}
492
493impl<const N: usize> Buf<N> {
494    #[inline(always)]
495    fn finish(&mut self, mut w: impl Write) -> io::Result<()> {
496        debug_assert_eq!(self.pos, N);
497        w.write_all(&self.buf)
498    }
499}
500
501#[test]
502fn error_cast() {
503    let _ : Box<dyn error::Error> = EncodingError::from(EncodingFormatError::MissingColorPalette).into();
504}