1 | //! Decoding of lossless WebP images |
2 | //! |
3 | //! [Lossless spec](https://developers.google.com/speed/webp/docs/webp_lossless_bitstream_specification) |
4 | |
5 | use std::io::BufRead; |
6 | use std::mem; |
7 | |
8 | use crate::decoder::DecodingError; |
9 | use crate::lossless_transform::{ |
10 | apply_color_indexing_transform, apply_color_transform, apply_predictor_transform, |
11 | apply_subtract_green_transform, |
12 | }; |
13 | |
14 | use super::huffman::HuffmanTree; |
15 | use super::lossless_transform::TransformType; |
16 | |
17 | const CODE_LENGTH_CODES: usize = 19; |
18 | const CODE_LENGTH_CODE_ORDER: [usize; CODE_LENGTH_CODES] = [ |
19 | 17, 18, 0, 1, 2, 3, 4, 5, 16, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, |
20 | ]; |
21 | |
22 | #[rustfmt::skip] |
23 | const DISTANCE_MAP: [(i8, i8); 120] = [ |
24 | (0, 1), (1, 0), (1, 1), (-1, 1), (0, 2), (2, 0), (1, 2), (-1, 2), |
25 | (2, 1), (-2, 1), (2, 2), (-2, 2), (0, 3), (3, 0), (1, 3), (-1, 3), |
26 | (3, 1), (-3, 1), (2, 3), (-2, 3), (3, 2), (-3, 2), (0, 4), (4, 0), |
27 | (1, 4), (-1, 4), (4, 1), (-4, 1), (3, 3), (-3, 3), (2, 4), (-2, 4), |
28 | (4, 2), (-4, 2), (0, 5), (3, 4), (-3, 4), (4, 3), (-4, 3), (5, 0), |
29 | (1, 5), (-1, 5), (5, 1), (-5, 1), (2, 5), (-2, 5), (5, 2), (-5, 2), |
30 | (4, 4), (-4, 4), (3, 5), (-3, 5), (5, 3), (-5, 3), (0, 6), (6, 0), |
31 | (1, 6), (-1, 6), (6, 1), (-6, 1), (2, 6), (-2, 6), (6, 2), (-6, 2), |
32 | (4, 5), (-4, 5), (5, 4), (-5, 4), (3, 6), (-3, 6), (6, 3), (-6, 3), |
33 | (0, 7), (7, 0), (1, 7), (-1, 7), (5, 5), (-5, 5), (7, 1), (-7, 1), |
34 | (4, 6), (-4, 6), (6, 4), (-6, 4), (2, 7), (-2, 7), (7, 2), (-7, 2), |
35 | (3, 7), (-3, 7), (7, 3), (-7, 3), (5, 6), (-5, 6), (6, 5), (-6, 5), |
36 | (8, 0), (4, 7), (-4, 7), (7, 4), (-7, 4), (8, 1), (8, 2), (6, 6), |
37 | (-6, 6), (8, 3), (5, 7), (-5, 7), (7, 5), (-7, 5), (8, 4), (6, 7), |
38 | (-6, 7), (7, 6), (-7, 6), (8, 5), (7, 7), (-7, 7), (8, 6), (8, 7) |
39 | ]; |
40 | |
41 | const GREEN: usize = 0; |
42 | const RED: usize = 1; |
43 | const BLUE: usize = 2; |
44 | const ALPHA: usize = 3; |
45 | const DIST: usize = 4; |
46 | |
47 | const HUFFMAN_CODES_PER_META_CODE: usize = 5; |
48 | |
49 | type HuffmanCodeGroup = [HuffmanTree; HUFFMAN_CODES_PER_META_CODE]; |
50 | |
51 | const ALPHABET_SIZE: [u16; HUFFMAN_CODES_PER_META_CODE] = [256 + 24, 256, 256, 256, 40]; |
52 | |
53 | #[inline ] |
54 | pub(crate) fn subsample_size(size: u16, bits: u8) -> u16 { |
55 | ((u32::from(size) + (1u32 << bits) - 1) >> bits) |
56 | .try_into() |
57 | .unwrap() |
58 | } |
59 | |
60 | const NUM_TRANSFORM_TYPES: usize = 4; |
61 | |
62 | //Decodes lossless WebP images |
63 | #[derive (Debug)] |
64 | pub(crate) struct LosslessDecoder<R> { |
65 | bit_reader: BitReader<R>, |
66 | transforms: [Option<TransformType>; NUM_TRANSFORM_TYPES], |
67 | transform_order: Vec<u8>, |
68 | width: u16, |
69 | height: u16, |
70 | } |
71 | |
72 | impl<R: BufRead> LosslessDecoder<R> { |
73 | /// Create a new decoder |
74 | pub(crate) const fn new(r: R) -> Self { |
75 | Self { |
76 | bit_reader: BitReader::new(r), |
77 | transforms: [None, None, None, None], |
78 | transform_order: Vec::new(), |
79 | width: 0, |
80 | height: 0, |
81 | } |
82 | } |
83 | |
84 | /// Decodes a frame. |
85 | /// |
86 | /// In an alpha chunk the width and height are not included in the header, so they should be |
87 | /// provided by setting the `implicit_dimensions` argument. Otherwise that argument should be |
88 | /// `None` and the frame dimensions will be determined by reading the VP8L header. |
89 | pub(crate) fn decode_frame( |
90 | &mut self, |
91 | width: u32, |
92 | height: u32, |
93 | implicit_dimensions: bool, |
94 | buf: &mut [u8], |
95 | ) -> Result<(), DecodingError> { |
96 | if implicit_dimensions { |
97 | self.width = width as u16; |
98 | self.height = height as u16; |
99 | } else { |
100 | let signature = self.bit_reader.read_bits::<u8>(8)?; |
101 | if signature != 0x2f { |
102 | return Err(DecodingError::LosslessSignatureInvalid(signature)); |
103 | } |
104 | |
105 | self.width = self.bit_reader.read_bits::<u16>(14)? + 1; |
106 | self.height = self.bit_reader.read_bits::<u16>(14)? + 1; |
107 | if u32::from(self.width) != width || u32::from(self.height) != height { |
108 | return Err(DecodingError::InconsistentImageSizes); |
109 | } |
110 | |
111 | let _alpha_used = self.bit_reader.read_bits::<u8>(1)?; |
112 | let version_num = self.bit_reader.read_bits::<u8>(3)?; |
113 | if version_num != 0 { |
114 | return Err(DecodingError::VersionNumberInvalid(version_num)); |
115 | } |
116 | } |
117 | |
118 | let transformed_width = self.read_transforms()?; |
119 | let transformed_size = usize::from(transformed_width) * usize::from(self.height) * 4; |
120 | self.decode_image_stream( |
121 | transformed_width, |
122 | self.height, |
123 | true, |
124 | &mut buf[..transformed_size], |
125 | )?; |
126 | |
127 | let mut image_size = transformed_size; |
128 | let mut width = transformed_width; |
129 | for &trans_index in self.transform_order.iter().rev() { |
130 | let transform = self.transforms[usize::from(trans_index)].as_ref().unwrap(); |
131 | match transform { |
132 | TransformType::PredictorTransform { |
133 | size_bits, |
134 | predictor_data, |
135 | } => apply_predictor_transform( |
136 | &mut buf[..image_size], |
137 | width, |
138 | self.height, |
139 | *size_bits, |
140 | predictor_data, |
141 | )?, |
142 | TransformType::ColorTransform { |
143 | size_bits, |
144 | transform_data, |
145 | } => { |
146 | apply_color_transform( |
147 | &mut buf[..image_size], |
148 | width, |
149 | *size_bits, |
150 | transform_data, |
151 | ); |
152 | } |
153 | TransformType::SubtractGreen => { |
154 | apply_subtract_green_transform(&mut buf[..image_size]); |
155 | } |
156 | TransformType::ColorIndexingTransform { |
157 | table_size, |
158 | table_data, |
159 | } => { |
160 | width = self.width; |
161 | image_size = usize::from(width) * usize::from(self.height) * 4; |
162 | apply_color_indexing_transform( |
163 | buf, |
164 | width, |
165 | self.height, |
166 | *table_size, |
167 | table_data, |
168 | ); |
169 | } |
170 | } |
171 | } |
172 | |
173 | Ok(()) |
174 | } |
175 | |
176 | /// Reads Image data from the bitstream |
177 | /// |
178 | /// Can be in any of the 5 roles described in the Specification. ARGB Image role has different |
179 | /// behaviour to the other 4. xsize and ysize describe the size of the blocks where each block |
180 | /// has its own entropy code |
181 | fn decode_image_stream( |
182 | &mut self, |
183 | xsize: u16, |
184 | ysize: u16, |
185 | is_argb_img: bool, |
186 | data: &mut [u8], |
187 | ) -> Result<(), DecodingError> { |
188 | let color_cache_bits = self.read_color_cache()?; |
189 | let color_cache = color_cache_bits.map(|bits| ColorCache { |
190 | color_cache_bits: bits, |
191 | color_cache: vec![[0; 4]; 1 << bits], |
192 | }); |
193 | |
194 | let huffman_info = self.read_huffman_codes(is_argb_img, xsize, ysize, color_cache)?; |
195 | self.decode_image_data(xsize, ysize, huffman_info, data) |
196 | } |
197 | |
198 | /// Reads transforms and their data from the bitstream |
199 | fn read_transforms(&mut self) -> Result<u16, DecodingError> { |
200 | let mut xsize = self.width; |
201 | |
202 | while self.bit_reader.read_bits::<u8>(1)? == 1 { |
203 | let transform_type_val = self.bit_reader.read_bits::<u8>(2)?; |
204 | |
205 | if self.transforms[usize::from(transform_type_val)].is_some() { |
206 | //can only have one of each transform, error |
207 | return Err(DecodingError::TransformError); |
208 | } |
209 | |
210 | self.transform_order.push(transform_type_val); |
211 | |
212 | let transform_type = match transform_type_val { |
213 | 0 => { |
214 | //predictor |
215 | |
216 | let size_bits = self.bit_reader.read_bits::<u8>(3)? + 2; |
217 | |
218 | let block_xsize = subsample_size(xsize, size_bits); |
219 | let block_ysize = subsample_size(self.height, size_bits); |
220 | |
221 | let mut predictor_data = |
222 | vec![0; usize::from(block_xsize) * usize::from(block_ysize) * 4]; |
223 | self.decode_image_stream(block_xsize, block_ysize, false, &mut predictor_data)?; |
224 | |
225 | TransformType::PredictorTransform { |
226 | size_bits, |
227 | predictor_data, |
228 | } |
229 | } |
230 | 1 => { |
231 | //color transform |
232 | |
233 | let size_bits = self.bit_reader.read_bits::<u8>(3)? + 2; |
234 | |
235 | let block_xsize = subsample_size(xsize, size_bits); |
236 | let block_ysize = subsample_size(self.height, size_bits); |
237 | |
238 | let mut transform_data = |
239 | vec![0; usize::from(block_xsize) * usize::from(block_ysize) * 4]; |
240 | self.decode_image_stream(block_xsize, block_ysize, false, &mut transform_data)?; |
241 | |
242 | TransformType::ColorTransform { |
243 | size_bits, |
244 | transform_data, |
245 | } |
246 | } |
247 | 2 => { |
248 | //subtract green |
249 | |
250 | TransformType::SubtractGreen |
251 | } |
252 | 3 => { |
253 | let color_table_size = self.bit_reader.read_bits::<u16>(8)? + 1; |
254 | |
255 | let mut color_map = vec![0; usize::from(color_table_size) * 4]; |
256 | self.decode_image_stream(color_table_size, 1, false, &mut color_map)?; |
257 | |
258 | let bits = if color_table_size <= 2 { |
259 | 3 |
260 | } else if color_table_size <= 4 { |
261 | 2 |
262 | } else if color_table_size <= 16 { |
263 | 1 |
264 | } else { |
265 | 0 |
266 | }; |
267 | xsize = subsample_size(xsize, bits); |
268 | |
269 | Self::adjust_color_map(&mut color_map); |
270 | |
271 | TransformType::ColorIndexingTransform { |
272 | table_size: color_table_size, |
273 | table_data: color_map, |
274 | } |
275 | } |
276 | _ => unreachable!(), |
277 | }; |
278 | |
279 | self.transforms[usize::from(transform_type_val)] = Some(transform_type); |
280 | } |
281 | |
282 | Ok(xsize) |
283 | } |
284 | |
285 | /// Adjusts the color map since it's subtraction coded |
286 | fn adjust_color_map(color_map: &mut [u8]) { |
287 | for i in 4..color_map.len() { |
288 | color_map[i] = color_map[i].wrapping_add(color_map[i - 4]); |
289 | } |
290 | } |
291 | |
292 | /// Reads huffman codes associated with an image |
293 | fn read_huffman_codes( |
294 | &mut self, |
295 | read_meta: bool, |
296 | xsize: u16, |
297 | ysize: u16, |
298 | color_cache: Option<ColorCache>, |
299 | ) -> Result<HuffmanInfo, DecodingError> { |
300 | let mut num_huff_groups = 1u32; |
301 | |
302 | let mut huffman_bits = 0; |
303 | let mut huffman_xsize = 1; |
304 | let mut huffman_ysize = 1; |
305 | let mut entropy_image = Vec::new(); |
306 | |
307 | if read_meta && self.bit_reader.read_bits::<u8>(1)? == 1 { |
308 | //meta huffman codes |
309 | huffman_bits = self.bit_reader.read_bits::<u8>(3)? + 2; |
310 | huffman_xsize = subsample_size(xsize, huffman_bits); |
311 | huffman_ysize = subsample_size(ysize, huffman_bits); |
312 | |
313 | let mut data = vec![0; usize::from(huffman_xsize) * usize::from(huffman_ysize) * 4]; |
314 | self.decode_image_stream(huffman_xsize, huffman_ysize, false, &mut data)?; |
315 | |
316 | entropy_image = data |
317 | .chunks_exact(4) |
318 | .map(|pixel| { |
319 | let meta_huff_code = (u16::from(pixel[0]) << 8) | u16::from(pixel[1]); |
320 | if u32::from(meta_huff_code) >= num_huff_groups { |
321 | num_huff_groups = u32::from(meta_huff_code) + 1; |
322 | } |
323 | meta_huff_code |
324 | }) |
325 | .collect::<Vec<u16>>(); |
326 | } |
327 | |
328 | let mut hufftree_groups = Vec::new(); |
329 | |
330 | for _i in 0..num_huff_groups { |
331 | let mut group: HuffmanCodeGroup = Default::default(); |
332 | for j in 0..HUFFMAN_CODES_PER_META_CODE { |
333 | let mut alphabet_size = ALPHABET_SIZE[j]; |
334 | if j == 0 { |
335 | if let Some(color_cache) = color_cache.as_ref() { |
336 | alphabet_size += 1 << color_cache.color_cache_bits; |
337 | } |
338 | } |
339 | |
340 | let tree = self.read_huffman_code(alphabet_size)?; |
341 | group[j] = tree; |
342 | } |
343 | hufftree_groups.push(group); |
344 | } |
345 | |
346 | let huffman_mask = if huffman_bits == 0 { |
347 | !0 |
348 | } else { |
349 | (1 << huffman_bits) - 1 |
350 | }; |
351 | |
352 | let info = HuffmanInfo { |
353 | xsize: huffman_xsize, |
354 | _ysize: huffman_ysize, |
355 | color_cache, |
356 | image: entropy_image, |
357 | bits: huffman_bits, |
358 | mask: huffman_mask, |
359 | huffman_code_groups: hufftree_groups, |
360 | }; |
361 | |
362 | Ok(info) |
363 | } |
364 | |
365 | /// Decodes and returns a single huffman tree |
366 | fn read_huffman_code(&mut self, alphabet_size: u16) -> Result<HuffmanTree, DecodingError> { |
367 | let simple = self.bit_reader.read_bits::<u8>(1)? == 1; |
368 | |
369 | if simple { |
370 | let num_symbols = self.bit_reader.read_bits::<u8>(1)? + 1; |
371 | |
372 | let is_first_8bits = self.bit_reader.read_bits::<u8>(1)?; |
373 | let zero_symbol = self.bit_reader.read_bits::<u16>(1 + 7 * is_first_8bits)?; |
374 | |
375 | if zero_symbol >= alphabet_size { |
376 | return Err(DecodingError::BitStreamError); |
377 | } |
378 | |
379 | if num_symbols == 1 { |
380 | Ok(HuffmanTree::build_single_node(zero_symbol)) |
381 | } else { |
382 | let one_symbol = self.bit_reader.read_bits::<u16>(8)?; |
383 | if one_symbol >= alphabet_size { |
384 | return Err(DecodingError::BitStreamError); |
385 | } |
386 | Ok(HuffmanTree::build_two_node(zero_symbol, one_symbol)) |
387 | } |
388 | } else { |
389 | let mut code_length_code_lengths = vec![0; CODE_LENGTH_CODES]; |
390 | |
391 | let num_code_lengths = 4 + self.bit_reader.read_bits::<usize>(4)?; |
392 | for i in 0..num_code_lengths { |
393 | code_length_code_lengths[CODE_LENGTH_CODE_ORDER[i]] = |
394 | self.bit_reader.read_bits(3)?; |
395 | } |
396 | |
397 | let new_code_lengths = |
398 | self.read_huffman_code_lengths(code_length_code_lengths, alphabet_size)?; |
399 | |
400 | HuffmanTree::build_implicit(new_code_lengths) |
401 | } |
402 | } |
403 | |
404 | /// Reads huffman code lengths |
405 | fn read_huffman_code_lengths( |
406 | &mut self, |
407 | code_length_code_lengths: Vec<u16>, |
408 | num_symbols: u16, |
409 | ) -> Result<Vec<u16>, DecodingError> { |
410 | let table = HuffmanTree::build_implicit(code_length_code_lengths)?; |
411 | |
412 | let mut max_symbol = if self.bit_reader.read_bits::<u8>(1)? == 1 { |
413 | let length_nbits = 2 + 2 * self.bit_reader.read_bits::<u8>(3)?; |
414 | let max_minus_two = self.bit_reader.read_bits::<u16>(length_nbits)?; |
415 | if max_minus_two > num_symbols - 2 { |
416 | return Err(DecodingError::BitStreamError); |
417 | } |
418 | 2 + max_minus_two |
419 | } else { |
420 | num_symbols |
421 | }; |
422 | |
423 | let mut code_lengths = vec![0; usize::from(num_symbols)]; |
424 | let mut prev_code_len = 8; //default code length |
425 | |
426 | let mut symbol = 0; |
427 | while symbol < num_symbols { |
428 | if max_symbol == 0 { |
429 | break; |
430 | } |
431 | max_symbol -= 1; |
432 | |
433 | self.bit_reader.fill()?; |
434 | let code_len = table.read_symbol(&mut self.bit_reader)?; |
435 | |
436 | if code_len < 16 { |
437 | code_lengths[usize::from(symbol)] = code_len; |
438 | symbol += 1; |
439 | if code_len != 0 { |
440 | prev_code_len = code_len; |
441 | } |
442 | } else { |
443 | let use_prev = code_len == 16; |
444 | let slot = code_len - 16; |
445 | let extra_bits = match slot { |
446 | 0 => 2, |
447 | 1 => 3, |
448 | 2 => 7, |
449 | _ => return Err(DecodingError::BitStreamError), |
450 | }; |
451 | let repeat_offset = match slot { |
452 | 0 | 1 => 3, |
453 | 2 => 11, |
454 | _ => return Err(DecodingError::BitStreamError), |
455 | }; |
456 | |
457 | let mut repeat = self.bit_reader.read_bits::<u16>(extra_bits)? + repeat_offset; |
458 | |
459 | if symbol + repeat > num_symbols { |
460 | return Err(DecodingError::BitStreamError); |
461 | } |
462 | |
463 | let length = if use_prev { prev_code_len } else { 0 }; |
464 | while repeat > 0 { |
465 | repeat -= 1; |
466 | code_lengths[usize::from(symbol)] = length; |
467 | symbol += 1; |
468 | } |
469 | } |
470 | } |
471 | |
472 | Ok(code_lengths) |
473 | } |
474 | |
475 | /// Decodes the image data using the huffman trees and either of the 3 methods of decoding |
476 | fn decode_image_data( |
477 | &mut self, |
478 | width: u16, |
479 | height: u16, |
480 | mut huffman_info: HuffmanInfo, |
481 | data: &mut [u8], |
482 | ) -> Result<(), DecodingError> { |
483 | let num_values = usize::from(width) * usize::from(height); |
484 | |
485 | let huff_index = huffman_info.get_huff_index(0, 0); |
486 | let mut tree = &huffman_info.huffman_code_groups[huff_index]; |
487 | let mut index = 0; |
488 | |
489 | let mut next_block_start = 0; |
490 | while index < num_values { |
491 | self.bit_reader.fill()?; |
492 | |
493 | if index >= next_block_start { |
494 | let x = index % usize::from(width); |
495 | let y = index / usize::from(width); |
496 | next_block_start = (x | usize::from(huffman_info.mask)).min(usize::from(width - 1)) |
497 | + y * usize::from(width) |
498 | + 1; |
499 | |
500 | let huff_index = huffman_info.get_huff_index(x as u16, y as u16); |
501 | tree = &huffman_info.huffman_code_groups[huff_index]; |
502 | |
503 | // Fast path: If all the codes each contain only a single |
504 | // symbol, then the pixel data isn't written to the bitstream |
505 | // and we can just fill the output buffer with the symbol |
506 | // directly. |
507 | if tree[..4].iter().all(|t| t.is_single_node()) { |
508 | let code = tree[GREEN].read_symbol(&mut self.bit_reader)?; |
509 | if code < 256 { |
510 | let n = if huffman_info.bits == 0 { |
511 | num_values |
512 | } else { |
513 | next_block_start - index |
514 | }; |
515 | |
516 | let red = tree[RED].read_symbol(&mut self.bit_reader)?; |
517 | let blue = tree[BLUE].read_symbol(&mut self.bit_reader)?; |
518 | let alpha = tree[ALPHA].read_symbol(&mut self.bit_reader)?; |
519 | let value = [red as u8, code as u8, blue as u8, alpha as u8]; |
520 | |
521 | for i in 0..n { |
522 | data[index * 4 + i * 4..][..4].copy_from_slice(&value); |
523 | } |
524 | |
525 | if let Some(color_cache) = huffman_info.color_cache.as_mut() { |
526 | color_cache.insert(value); |
527 | } |
528 | |
529 | index += n; |
530 | continue; |
531 | } |
532 | } |
533 | } |
534 | |
535 | let code = tree[GREEN].read_symbol(&mut self.bit_reader)?; |
536 | |
537 | //check code |
538 | if code < 256 { |
539 | //literal, so just use huffman codes and read as argb |
540 | let green = code as u8; |
541 | let red = tree[RED].read_symbol(&mut self.bit_reader)? as u8; |
542 | let blue = tree[BLUE].read_symbol(&mut self.bit_reader)? as u8; |
543 | if self.bit_reader.nbits < 15 { |
544 | self.bit_reader.fill()?; |
545 | } |
546 | let alpha = tree[ALPHA].read_symbol(&mut self.bit_reader)? as u8; |
547 | |
548 | data[index * 4] = red; |
549 | data[index * 4 + 1] = green; |
550 | data[index * 4 + 2] = blue; |
551 | data[index * 4 + 3] = alpha; |
552 | |
553 | if let Some(color_cache) = huffman_info.color_cache.as_mut() { |
554 | color_cache.insert([red, green, blue, alpha]); |
555 | } |
556 | index += 1; |
557 | } else if code < 256 + 24 { |
558 | //backward reference, so go back and use that to add image data |
559 | let length_symbol = code - 256; |
560 | let length = Self::get_copy_distance(&mut self.bit_reader, length_symbol)?; |
561 | |
562 | let dist_symbol = tree[DIST].read_symbol(&mut self.bit_reader)?; |
563 | let dist_code = Self::get_copy_distance(&mut self.bit_reader, dist_symbol)?; |
564 | let dist = Self::plane_code_to_distance(width, dist_code); |
565 | |
566 | if index < dist || num_values - index < length { |
567 | return Err(DecodingError::BitStreamError); |
568 | } |
569 | |
570 | if dist == 1 { |
571 | let value: [u8; 4] = data[(index - dist) * 4..][..4].try_into().unwrap(); |
572 | for i in 0..length { |
573 | data[index * 4 + i * 4..][..4].copy_from_slice(&value); |
574 | } |
575 | } else { |
576 | if index + length + 3 <= num_values { |
577 | let start = (index - dist) * 4; |
578 | data.copy_within(start..start + 16, index * 4); |
579 | |
580 | if length > 4 || dist < 4 { |
581 | for i in (0..length * 4).step_by((dist * 4).min(16)).skip(1) { |
582 | data.copy_within(start + i..start + i + 16, index * 4 + i); |
583 | } |
584 | } |
585 | } else { |
586 | for i in 0..length * 4 { |
587 | data[index * 4 + i] = data[index * 4 + i - dist * 4]; |
588 | } |
589 | } |
590 | |
591 | if let Some(color_cache) = huffman_info.color_cache.as_mut() { |
592 | for pixel in data[index * 4..][..length * 4].chunks_exact(4) { |
593 | color_cache.insert(pixel.try_into().unwrap()); |
594 | } |
595 | } |
596 | } |
597 | index += length; |
598 | } else { |
599 | //color cache, so use previously stored pixels to get this pixel |
600 | let color_cache = huffman_info |
601 | .color_cache |
602 | .as_mut() |
603 | .ok_or(DecodingError::BitStreamError)?; |
604 | let color = color_cache.lookup((code - 280).into()); |
605 | data[index * 4..][..4].copy_from_slice(&color); |
606 | index += 1; |
607 | |
608 | if index < next_block_start { |
609 | if let Some((bits, code)) = tree[GREEN].peek_symbol(&self.bit_reader) { |
610 | if code >= 280 { |
611 | self.bit_reader.consume(bits)?; |
612 | data[index * 4..][..4] |
613 | .copy_from_slice(&color_cache.lookup((code - 280).into())); |
614 | index += 1; |
615 | } |
616 | } |
617 | } |
618 | } |
619 | } |
620 | |
621 | Ok(()) |
622 | } |
623 | |
624 | /// Reads color cache data from the bitstream |
625 | fn read_color_cache(&mut self) -> Result<Option<u8>, DecodingError> { |
626 | if self.bit_reader.read_bits::<u8>(1)? == 1 { |
627 | let code_bits = self.bit_reader.read_bits::<u8>(4)?; |
628 | |
629 | if !(1..=11).contains(&code_bits) { |
630 | return Err(DecodingError::InvalidColorCacheBits(code_bits)); |
631 | } |
632 | |
633 | Ok(Some(code_bits)) |
634 | } else { |
635 | Ok(None) |
636 | } |
637 | } |
638 | |
639 | /// Gets the copy distance from the prefix code and bitstream |
640 | fn get_copy_distance( |
641 | bit_reader: &mut BitReader<R>, |
642 | prefix_code: u16, |
643 | ) -> Result<usize, DecodingError> { |
644 | if prefix_code < 4 { |
645 | return Ok(usize::from(prefix_code + 1)); |
646 | } |
647 | let extra_bits: u8 = ((prefix_code - 2) >> 1).try_into().unwrap(); |
648 | let offset = (2 + (usize::from(prefix_code) & 1)) << extra_bits; |
649 | |
650 | let bits = bit_reader.peek(extra_bits) as usize; |
651 | bit_reader.consume(extra_bits)?; |
652 | |
653 | Ok(offset + bits + 1) |
654 | } |
655 | |
656 | /// Gets distance to pixel |
657 | fn plane_code_to_distance(xsize: u16, plane_code: usize) -> usize { |
658 | if plane_code > 120 { |
659 | plane_code - 120 |
660 | } else { |
661 | let (xoffset, yoffset) = DISTANCE_MAP[plane_code - 1]; |
662 | |
663 | let dist = i32::from(xoffset) + i32::from(yoffset) * i32::from(xsize); |
664 | if dist < 1 { |
665 | return 1; |
666 | } |
667 | dist.try_into().unwrap() |
668 | } |
669 | } |
670 | } |
671 | |
672 | #[derive (Debug, Clone)] |
673 | struct HuffmanInfo { |
674 | xsize: u16, |
675 | _ysize: u16, |
676 | color_cache: Option<ColorCache>, |
677 | image: Vec<u16>, |
678 | bits: u8, |
679 | mask: u16, |
680 | huffman_code_groups: Vec<HuffmanCodeGroup>, |
681 | } |
682 | |
683 | impl HuffmanInfo { |
684 | fn get_huff_index(&self, x: u16, y: u16) -> usize { |
685 | if self.bits == 0 { |
686 | return 0; |
687 | } |
688 | let position: usize = |
689 | usize::from(y >> self.bits) * usize::from(self.xsize) + usize::from(x >> self.bits); |
690 | let meta_huff_code: usize = usize::from(self.image[position]); |
691 | meta_huff_code |
692 | } |
693 | } |
694 | |
695 | #[derive (Debug, Clone)] |
696 | struct ColorCache { |
697 | color_cache_bits: u8, |
698 | color_cache: Vec<[u8; 4]>, |
699 | } |
700 | |
701 | impl ColorCache { |
702 | #[inline (always)] |
703 | fn insert(&mut self, color: [u8; 4]) { |
704 | let [r: u8, g: u8, b: u8, a: u8] = color; |
705 | let color_u32: u32 = |
706 | (u32::from(r) << 16) | (u32::from(g) << 8) | (u32::from(b)) | (u32::from(a) << 24); |
707 | let index: u32 = (0x1e35a7bdu32.wrapping_mul(color_u32)) >> (32 - self.color_cache_bits); |
708 | self.color_cache[index as usize] = color; |
709 | } |
710 | |
711 | #[inline (always)] |
712 | fn lookup(&self, index: usize) -> [u8; 4] { |
713 | self.color_cache[index] |
714 | } |
715 | } |
716 | |
717 | #[derive (Debug, Clone)] |
718 | pub(crate) struct BitReader<R> { |
719 | reader: R, |
720 | buffer: u64, |
721 | nbits: u8, |
722 | } |
723 | |
724 | impl<R: BufRead> BitReader<R> { |
725 | const fn new(reader: R) -> Self { |
726 | Self { |
727 | reader, |
728 | buffer: 0, |
729 | nbits: 0, |
730 | } |
731 | } |
732 | |
733 | /// Fills the buffer with bits from the input stream. |
734 | /// |
735 | /// After this function, the internal buffer will contain 64-bits or have reached the end of |
736 | /// the input stream. |
737 | pub(crate) fn fill(&mut self) -> Result<(), DecodingError> { |
738 | debug_assert!(self.nbits < 64); |
739 | |
740 | let mut buf = self.reader.fill_buf()?; |
741 | if buf.len() >= 8 { |
742 | let lookahead = u64::from_le_bytes(buf[..8].try_into().unwrap()); |
743 | self.reader.consume(usize::from((63 - self.nbits) / 8)); |
744 | self.buffer |= lookahead << self.nbits; |
745 | self.nbits |= 56; |
746 | } else { |
747 | while !buf.is_empty() && self.nbits < 56 { |
748 | self.buffer |= u64::from(buf[0]) << self.nbits; |
749 | self.nbits += 8; |
750 | self.reader.consume(1); |
751 | buf = self.reader.fill_buf()?; |
752 | } |
753 | } |
754 | |
755 | Ok(()) |
756 | } |
757 | |
758 | /// Peeks at the next `num` bits in the buffer. |
759 | pub(crate) const fn peek(&self, num: u8) -> u64 { |
760 | self.buffer & ((1 << num) - 1) |
761 | } |
762 | |
763 | /// Peeks at the full buffer. |
764 | pub(crate) const fn peek_full(&self) -> u64 { |
765 | self.buffer |
766 | } |
767 | |
768 | /// Consumes `num` bits from the buffer returning an error if there are not enough bits. |
769 | pub(crate) fn consume(&mut self, num: u8) -> Result<(), DecodingError> { |
770 | if self.nbits < num { |
771 | return Err(DecodingError::BitStreamError); |
772 | } |
773 | |
774 | self.buffer >>= num; |
775 | self.nbits -= num; |
776 | Ok(()) |
777 | } |
778 | |
779 | /// Convenience function to read a number of bits and convert them to a type. |
780 | pub(crate) fn read_bits<T: TryFrom<u32>>(&mut self, num: u8) -> Result<T, DecodingError> { |
781 | debug_assert!(num as usize <= 8 * mem::size_of::<T>()); |
782 | debug_assert!(num <= 32); |
783 | |
784 | if self.nbits < num { |
785 | self.fill()?; |
786 | } |
787 | let value = self.peek(num) as u32; |
788 | self.consume(num)?; |
789 | |
790 | value.try_into().map_err(|_| { |
791 | debug_assert!(false, "Value too large to fit in type" ); |
792 | DecodingError::BitStreamError |
793 | }) |
794 | } |
795 | } |
796 | |
797 | #[cfg (test)] |
798 | mod test { |
799 | |
800 | use std::io::Cursor; |
801 | |
802 | use super::BitReader; |
803 | |
804 | #[test ] |
805 | fn bit_read_test() { |
806 | //10011100 01000001 11100001 |
807 | let mut bit_reader = BitReader::new(Cursor::new(vec![0x9C, 0x41, 0xE1])); |
808 | |
809 | assert_eq!(bit_reader.read_bits::<u8>(3).unwrap(), 4); //100 |
810 | assert_eq!(bit_reader.read_bits::<u8>(2).unwrap(), 3); //11 |
811 | assert_eq!(bit_reader.read_bits::<u8>(6).unwrap(), 12); //001100 |
812 | assert_eq!(bit_reader.read_bits::<u16>(10).unwrap(), 40); //0000101000 |
813 | assert_eq!(bit_reader.read_bits::<u8>(3).unwrap(), 7); //111 |
814 | } |
815 | |
816 | #[test ] |
817 | fn bit_read_error_test() { |
818 | //01101010 |
819 | let mut bit_reader = BitReader::new(Cursor::new(vec![0x6A])); |
820 | |
821 | assert_eq!(bit_reader.read_bits::<u8>(3).unwrap(), 2); //010 |
822 | assert_eq!(bit_reader.read_bits::<u8>(5).unwrap(), 13); //01101 |
823 | assert!(bit_reader.read_bits::<u8>(4).is_err()); //error |
824 | } |
825 | } |
826 | |