1use 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#[derive(Debug)]
15#[non_exhaustive]
16pub enum EncodingFormatError {
17 TooManyColors,
19 MissingColorPalette,
21 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)]
38pub enum EncodingError {
40 Format(EncodingFormatError),
42 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#[derive(Copy, Clone, Debug, PartialEq, Eq)]
82pub enum Repeat {
83 Finite(u16),
85 Infinite,
87}
88
89impl Default for Repeat {
90 fn default() -> Self {
91 Self::Finite(0)
92 }
93}
94
95#[non_exhaustive]
97pub enum ExtensionData {
98 Control {
100 flags: u8,
102 delay: u16,
104 trns: u8,
106 },
107 Repetitions(Repeat),
109}
110
111impl ExtensionData {
112 #[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 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 pub fn set_repeat(&mut self, repeat: Repeat) -> Result<(), EncodingError> {
147 self.write_extension(ExtensionData::Repetitions(repeat))
148 }
149
150 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 flags |= table_size;
158 flags |= table_size << 4; self.write_screen_desc(flags)?;
162 Self::write_color_table(self.writer()?, palette, padding)?;
163 Ok(self)
164 }
165
166 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 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 for _ in 0..padding {
247 writer.write_all(&[0, 0, 0])?;
248 }
249 Ok(())
250 }
251
252 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 pub fn write_extension(&mut self, extension: ExtensionData) -> Result<(), EncodingError> {
267 use self::ExtensionData::*;
268 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 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 pub fn write_lzw_pre_encoded_frame(&mut self, frame: &Frame<'_>) -> Result<(), EncodingError> {
325 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 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)?; tmp.write_le(0u8)?; tmp.write_le(0u8)?; tmp.finish(self.writer()?)
347 }
348
349 pub fn get_ref(&self) -> &W {
351 self.w.as_ref().unwrap()
352 }
353
354 pub fn get_mut(&mut self) -> &mut W {
358 self.w.as_mut().unwrap()
359 }
360
361 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 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
378fn 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 if byte > 127 {
388 break;
389 }
390 }
391 }
392 let palette_min_len = max_byte as u32 + 1;
393 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 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
413pub 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
438fn 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}