| 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 | |