1use std::borrow::Cow;
2use std::io;
3use std::mem;
4use std::iter;
5use crate::common::Frame;
6use crate::MemoryLimit;
7
8use super::decoder::{
9 PLTE_CHANNELS, DecodingError, OutputBuffer
10};
11
12pub(crate) const N_CHANNELS: usize = 4;
13
14/// Output mode for the image data
15#[derive(Clone, Copy, Debug, PartialEq)]
16#[repr(u8)]
17pub enum ColorOutput {
18 /// The decoder expands the image data to 32bit RGBA.
19 /// This affects:
20 ///
21 /// - The buffer buffer of the `Frame` returned by [`Decoder::read_next_frame`].
22 /// - `Decoder::fill_buffer`, `Decoder::buffer_size` and `Decoder::line_length`.
23 RGBA = 0,
24 /// The decoder returns the raw indexed data.
25 Indexed = 1,
26}
27
28pub(crate) type FillBufferCallback<'a> = &'a mut dyn FnMut(&mut OutputBuffer<'_>) -> Result<usize, DecodingError>;
29
30/// Deinterlaces and expands to RGBA if needed
31pub(crate) struct PixelConverter {
32 memory_limit: MemoryLimit,
33 color_output: ColorOutput,
34 buffer: Vec<u8>,
35 global_palette: Option<Vec<u8>>,
36}
37
38impl PixelConverter {
39 pub(crate) fn new(color_output: ColorOutput, memory_limit: MemoryLimit) -> Self {
40 Self {
41 memory_limit,
42 color_output,
43 buffer: Vec::new(),
44 global_palette: None,
45 }
46 }
47
48 pub(crate) fn check_buffer_size(&mut self, frame: &Frame<'_>) -> Result<usize, DecodingError> {
49 let pixel_bytes = self.memory_limit
50 .buffer_size(self.color_output, frame.width, frame.height)
51 .ok_or_else(|| io::Error::new(io::ErrorKind::OutOfMemory, "image is too large"))?;
52
53 debug_assert_eq!(
54 pixel_bytes, self.buffer_size(frame).unwrap(),
55 "Checked computation diverges from required buffer size"
56 );
57 Ok(pixel_bytes)
58 }
59
60 #[inline]
61 pub(crate) fn read_frame(&mut self, frame: &mut Frame<'_>, data_callback: FillBufferCallback<'_>) -> Result<(), DecodingError> {
62 let pixel_bytes = self.check_buffer_size(frame)?;
63 let mut vec = match mem::replace(&mut frame.buffer, Cow::Borrowed(&[])) {
64 // reuse buffer if possible without reallocating
65 Cow::Owned(mut vec) if vec.capacity() >= pixel_bytes => {
66 vec.resize(pixel_bytes, 0);
67 vec
68 },
69 // resizing would realloc anyway, and 0-init is faster than a copy
70 _ => vec![0; pixel_bytes],
71 };
72 self.read_into_buffer(frame, &mut vec, data_callback)?;
73 frame.buffer = Cow::Owned(vec);
74 frame.interlaced = false;
75 Ok(())
76 }
77
78 #[inline]
79 pub(crate) fn buffer_size(&self, frame: &Frame<'_>) -> Option<usize> {
80 self.line_length(frame).checked_mul(frame.height as usize)
81 }
82
83 #[inline]
84 pub(crate) fn line_length(&self, frame: &Frame<'_>) -> usize {
85 use self::ColorOutput::*;
86 match self.color_output {
87 RGBA => frame.width as usize * N_CHANNELS,
88 Indexed => frame.width as usize,
89 }
90 }
91
92 /// Use `read_into_buffer` to deinterlace
93 #[inline(never)]
94 pub(crate) fn fill_buffer(&mut self, current_frame: &Frame<'_>, mut buf: &mut [u8], data_callback: FillBufferCallback<'_>) -> Result<bool, DecodingError> {
95 loop {
96 let decode_into = match self.color_output {
97 // When decoding indexed data, LZW can write the pixels directly
98 ColorOutput::Indexed => &mut buf[..],
99 // When decoding RGBA, the pixel data will be expanded by a factor of 4,
100 // and it's simpler to decode indexed pixels to another buffer first
101 ColorOutput::RGBA => {
102 let buffer_size = buf.len() / N_CHANNELS;
103 if buffer_size == 0 {
104 return Err(DecodingError::format("odd-sized buffer"));
105 }
106 if self.buffer.len() < buffer_size {
107 self.buffer.resize(buffer_size, 0);
108 }
109 &mut self.buffer[..buffer_size]
110 }
111 };
112 match data_callback(&mut OutputBuffer::Slice(decode_into))? {
113 0 => return Ok(false),
114 bytes_decoded => {
115 match self.color_output {
116 ColorOutput::RGBA => {
117 let transparent = current_frame.transparent;
118 let palette: &[u8] = current_frame.palette.as_deref()
119 .or(self.global_palette.as_deref())
120 .unwrap_or_default(); // next_frame_info already checked it won't happen
121
122 let (pixels, rest) = buf.split_at_mut(bytes_decoded * N_CHANNELS);
123 buf = rest;
124
125 for (rgba, idx) in pixels.chunks_exact_mut(N_CHANNELS).zip(self.buffer.iter().copied().take(bytes_decoded)) {
126 let plte_offset = PLTE_CHANNELS * idx as usize;
127 if let Some(colors) = palette.get(plte_offset..plte_offset+PLTE_CHANNELS) {
128 rgba[0] = colors[0];
129 rgba[1] = colors[1];
130 rgba[2] = colors[2];
131 rgba[3] = if let Some(t) = transparent {
132 if t == idx { 0x00 } else { 0xFF }
133 } else {
134 0xFF
135 };
136 }
137 }
138 },
139 ColorOutput::Indexed => {
140 buf = &mut buf[bytes_decoded..];
141 }
142 }
143 if buf.is_empty() {
144 return Ok(true);
145 }
146 },
147 }
148 }
149 }
150
151 pub(crate) fn global_palette(&self) -> Option<&[u8]> {
152 self.global_palette.as_deref()
153 }
154
155 pub(crate) fn set_global_palette(&mut self, palette: Vec<u8>) {
156 self.global_palette = if !palette.is_empty() {
157 Some(palette)
158 } else {
159 None
160 };
161 }
162
163 /// Applies deinterlacing
164 ///
165 /// Set `frame.interlaced = false` afterwards if you're putting the buffer back into the `Frame`
166 pub(crate) fn read_into_buffer(&mut self, frame: &Frame<'_>, buf: &mut [u8], data_callback: FillBufferCallback<'_>) -> Result<(), DecodingError> {
167 if frame.interlaced {
168 let width = self.line_length(frame);
169 for row in (InterlaceIterator { len: frame.height, next: 0, pass: 0 }) {
170 // this can't overflow 32-bit, because row never equals (maximum) height
171 let start = row * width;
172 // Handle a too-small buffer and 32-bit usize overflow without panicking
173 let line = buf.get_mut(start..).and_then(|b| b.get_mut(..width))
174 .ok_or_else(|| DecodingError::format("buffer too small"))?;
175 if !self.fill_buffer(frame, line, data_callback)? {
176 return Err(DecodingError::format("image truncated"));
177 }
178 }
179 } else {
180 let buf = self.buffer_size(frame).and_then(|buffer_size| buf.get_mut(..buffer_size))
181 .ok_or_else(|| DecodingError::format("buffer too small"))?;
182 if !self.fill_buffer(frame, buf, data_callback)? {
183 return Err(DecodingError::format("image truncated"));
184 }
185 };
186 Ok(())
187 }
188}
189
190struct InterlaceIterator {
191 len: u16,
192 next: usize,
193 pass: usize,
194}
195
196impl iter::Iterator for InterlaceIterator {
197 type Item = usize;
198
199 #[inline]
200 fn next(&mut self) -> Option<Self::Item> {
201 if self.len == 0 {
202 return None;
203 }
204 // although the pass never goes out of bounds thanks to len==0,
205 // the optimizer doesn't see it. get()? avoids costlier panicking code.
206 let mut next: usize = self.next + *[8, 8, 4, 2].get(self.pass)?;
207 while next >= self.len as usize {
208 debug_assert!(self.pass < 4);
209 next = *[4, 2, 1, 0].get(self.pass)?;
210 self.pass += 1;
211 }
212 mem::swap(&mut next, &mut self.next);
213 Some(next)
214 }
215}
216
217#[cfg(test)]
218mod test {
219 use super::InterlaceIterator;
220
221 #[test]
222 fn test_interlace_iterator() {
223 for &(len, expect) in &[
224 (0, &[][..]),
225 (1, &[0][..]),
226 (2, &[0, 1][..]),
227 (3, &[0, 2, 1][..]),
228 (4, &[0, 2, 1, 3][..]),
229 (5, &[0, 4, 2, 1, 3][..]),
230 (6, &[0, 4, 2, 1, 3, 5][..]),
231 (7, &[0, 4, 2, 6, 1, 3, 5][..]),
232 (8, &[0, 4, 2, 6, 1, 3, 5, 7][..]),
233 (9, &[0, 8, 4, 2, 6, 1, 3, 5, 7][..]),
234 (10, &[0, 8, 4, 2, 6, 1, 3, 5, 7, 9][..]),
235 (11, &[0, 8, 4, 2, 6, 10, 1, 3, 5, 7, 9][..]),
236 (12, &[0, 8, 4, 2, 6, 10, 1, 3, 5, 7, 9, 11][..]),
237 (13, &[0, 8, 4, 12, 2, 6, 10, 1, 3, 5, 7, 9, 11][..]),
238 (14, &[0, 8, 4, 12, 2, 6, 10, 1, 3, 5, 7, 9, 11, 13][..]),
239 (15, &[0, 8, 4, 12, 2, 6, 10, 14, 1, 3, 5, 7, 9, 11, 13][..]),
240 (16, &[0, 8, 4, 12, 2, 6, 10, 14, 1, 3, 5, 7, 9, 11, 13, 15][..]),
241 (17, &[0, 8, 16, 4, 12, 2, 6, 10, 14, 1, 3, 5, 7, 9, 11, 13, 15][..]),
242 ] {
243 let iter = InterlaceIterator { len, next: 0, pass: 0 };
244 let lines = iter.collect::<Vec<_>>();
245 assert_eq!(lines, expect);
246 }
247 }
248
249 #[test]
250 fn interlace_max() {
251 let iter = InterlaceIterator { len: 0xFFFF, next: 0, pass: 0 };
252 assert_eq!(65533, iter.last().unwrap());
253 }
254}
255