1use alloc::borrow::ToOwned;
2use alloc::vec;
3use alloc::vec::Vec;
4use core::iter;
5use std::io::Read;
6use crate::read_u8;
7use crate::error::{Error, Result};
8use crate::marker::Marker;
9use crate::parser::ScanInfo;
10
11const LUT_BITS: u8 = 8;
12
13#[derive(Debug)]
14pub struct HuffmanDecoder {
15 bits: u64,
16 num_bits: u8,
17 marker: Option<Marker>,
18}
19
20impl HuffmanDecoder {
21 pub fn new() -> HuffmanDecoder {
22 HuffmanDecoder {
23 bits: 0,
24 num_bits: 0,
25 marker: None,
26 }
27 }
28
29 // Section F.2.2.3
30 // Figure F.16
31 pub fn decode<R: Read>(&mut self, reader: &mut R, table: &HuffmanTable) -> Result<u8> {
32 if self.num_bits < 16 {
33 self.read_bits(reader)?;
34 }
35
36 let (value, size) = table.lut[self.peek_bits(LUT_BITS) as usize];
37
38 if size > 0 {
39 self.consume_bits(size);
40 Ok(value)
41 }
42 else {
43 let bits = self.peek_bits(16);
44
45 for i in LUT_BITS .. 16 {
46 let code = (bits >> (15 - i)) as i32;
47
48 if code <= table.maxcode[i as usize] {
49 self.consume_bits(i + 1);
50
51 let index = (code + table.delta[i as usize]) as usize;
52 return Ok(table.values[index]);
53 }
54 }
55
56 Err(Error::Format("failed to decode huffman code".to_owned()))
57 }
58 }
59
60 pub fn decode_fast_ac<R: Read>(&mut self, reader: &mut R, table: &HuffmanTable) -> Result<Option<(i16, u8)>> {
61 if let Some(ref ac_lut) = table.ac_lut {
62 if self.num_bits < LUT_BITS {
63 self.read_bits(reader)?;
64 }
65
66 let (value, run_size) = ac_lut[self.peek_bits(LUT_BITS) as usize];
67
68 if run_size != 0 {
69 let run = run_size >> 4;
70 let size = run_size & 0x0f;
71
72 self.consume_bits(size);
73 return Ok(Some((value, run)));
74 }
75 }
76
77 Ok(None)
78 }
79
80 #[inline]
81 pub fn get_bits<R: Read>(&mut self, reader: &mut R, count: u8) -> Result<u16> {
82 if self.num_bits < count {
83 self.read_bits(reader)?;
84 }
85
86 let bits = self.peek_bits(count);
87 self.consume_bits(count);
88
89 Ok(bits)
90 }
91
92 #[inline]
93 pub fn receive_extend<R: Read>(&mut self, reader: &mut R, count: u8) -> Result<i16> {
94 let value = self.get_bits(reader, count)?;
95 Ok(extend(value, count))
96 }
97
98 pub fn reset(&mut self) {
99 self.bits = 0;
100 self.num_bits = 0;
101 }
102
103 pub fn take_marker<R: Read>(&mut self, reader: &mut R) -> Result<Option<Marker>> {
104 self.read_bits(reader).map(|_| self.marker.take())
105 }
106
107 #[inline]
108 fn peek_bits(&mut self, count: u8) -> u16 {
109 debug_assert!(count <= 16);
110 debug_assert!(self.num_bits >= count);
111
112 ((self.bits >> (64 - count)) & ((1 << count) - 1)) as u16
113 }
114
115 #[inline]
116 fn consume_bits(&mut self, count: u8) {
117 debug_assert!(self.num_bits >= count);
118
119 self.bits <<= count as usize;
120 self.num_bits -= count;
121 }
122
123 fn read_bits<R: Read>(&mut self, reader: &mut R) -> Result<()> {
124 while self.num_bits <= 56 {
125 // Fill with zero bits if we have reached the end.
126 let byte = match self.marker {
127 Some(_) => 0,
128 None => read_u8(reader)?,
129 };
130
131 if byte == 0xFF {
132 let mut next_byte = read_u8(reader)?;
133
134 // Check for byte stuffing.
135 if next_byte != 0x00 {
136 // We seem to have reached the end of entropy-coded data and encountered a
137 // marker. Since we can't put data back into the reader, we have to continue
138 // reading to identify the marker so we can pass it on.
139
140 // Section B.1.1.2
141 // "Any marker may optionally be preceded by any number of fill bytes, which are bytes assigned code X’FF’."
142 while next_byte == 0xFF {
143 next_byte = read_u8(reader)?;
144 }
145
146 match next_byte {
147 0x00 => return Err(Error::Format("FF 00 found where marker was expected".to_owned())),
148 _ => self.marker = Some(Marker::from_u8(next_byte).unwrap()),
149 }
150
151 continue;
152 }
153 }
154
155 self.bits |= (byte as u64) << (56 - self.num_bits);
156 self.num_bits += 8;
157 }
158
159 Ok(())
160 }
161}
162
163// Section F.2.2.1
164// Figure F.12
165fn extend(value: u16, count: u8) -> i16 {
166 let vt: u16 = 1 << (count as u16 - 1);
167
168 if value < vt {
169 value as i16 + (-1 << count as i16) + 1
170 } else {
171 value as i16
172 }
173}
174
175#[derive(Clone, Copy, Debug, PartialEq)]
176pub enum HuffmanTableClass {
177 DC,
178 AC,
179}
180
181pub struct HuffmanTable {
182 values: Vec<u8>,
183 delta: [i32; 16],
184 maxcode: [i32; 16],
185
186 lut: [(u8, u8); 1 << LUT_BITS],
187 ac_lut: Option<[(i16, u8); 1 << LUT_BITS]>,
188}
189
190impl HuffmanTable {
191 pub fn new(bits: &[u8; 16], values: &[u8], class: HuffmanTableClass) -> Result<HuffmanTable> {
192 let (huffcode, huffsize) = derive_huffman_codes(bits)?;
193
194 // Section F.2.2.3
195 // Figure F.15
196 // delta[i] is set to VALPTR(I) - MINCODE(I)
197 let mut delta = [0i32; 16];
198 let mut maxcode = [-1i32; 16];
199 let mut j = 0;
200
201 for i in 0 .. 16 {
202 if bits[i] != 0 {
203 delta[i] = j as i32 - huffcode[j] as i32;
204 j += bits[i] as usize;
205 maxcode[i] = huffcode[j - 1] as i32;
206 }
207 }
208
209 // Build a lookup table for faster decoding.
210 let mut lut = [(0u8, 0u8); 1 << LUT_BITS];
211
212 for (i, &size) in huffsize.iter().enumerate().filter(|&(_, &size)| size <= LUT_BITS) {
213 let bits_remaining = LUT_BITS - size;
214 let start = (huffcode[i] << bits_remaining) as usize;
215
216 let val = (values[i], size);
217 for b in &mut lut[start..][..1 << bits_remaining] {
218 *b = val;
219 }
220 }
221
222 // Build a lookup table for small AC coefficients which both decodes the value and does the
223 // equivalent of receive_extend.
224 let ac_lut = match class {
225 HuffmanTableClass::DC => None,
226 HuffmanTableClass::AC => {
227 let mut table = [(0i16, 0u8); 1 << LUT_BITS];
228
229 for (i, &(value, size)) in lut.iter().enumerate() {
230 let run_length = value >> 4;
231 let magnitude_category = value & 0x0f;
232
233 if magnitude_category > 0 && size + magnitude_category <= LUT_BITS {
234 let unextended_ac_value = (((i << size) & ((1 << LUT_BITS) - 1)) >> (LUT_BITS - magnitude_category)) as u16;
235 let ac_value = extend(unextended_ac_value, magnitude_category);
236
237 table[i] = (ac_value, (run_length << 4) | (size + magnitude_category));
238 }
239 }
240
241 Some(table)
242 },
243 };
244
245 Ok(HuffmanTable {
246 values: values.to_vec(),
247 delta,
248 maxcode,
249 lut,
250 ac_lut,
251 })
252 }
253}
254
255// Section C.2
256fn derive_huffman_codes(bits: &[u8; 16]) -> Result<(Vec<u16>, Vec<u8>)> {
257 // Figure C.1
258 let huffsize = bits.iter()
259 .enumerate()
260 .fold(Vec::new(), |mut acc, (i, &value)| {
261 acc.extend(iter::repeat((i + 1) as u8).take(value as usize));
262 acc
263 });
264
265 // Figure C.2
266 let mut huffcode = vec![0u16; huffsize.len()];
267 let mut code_size = huffsize[0];
268 let mut code = 0u32;
269
270 for (i, &size) in huffsize.iter().enumerate() {
271 while code_size < size {
272 code <<= 1;
273 code_size += 1;
274 }
275
276 if code >= (1u32 << size) {
277 return Err(Error::Format("bad huffman code length".to_owned()));
278 }
279
280 huffcode[i] = code as u16;
281 code += 1;
282 }
283
284 Ok((huffcode, huffsize))
285}
286
287// https://www.loc.gov/preservation/digital/formats/fdd/fdd000063.shtml
288// "Avery Lee, writing in the rec.video.desktop newsgroup in 2001, commented that "MJPEG, or at
289// least the MJPEG in AVIs having the MJPG fourcc, is restricted JPEG with a fixed -- and
290// *omitted* -- Huffman table. The JPEG must be YCbCr colorspace, it must be 4:2:2, and it must
291// use basic Huffman encoding, not arithmetic or progressive.... You can indeed extract the
292// MJPEG frames and decode them with a regular JPEG decoder, but you have to prepend the DHT
293// segment to them, or else the decoder won't have any idea how to decompress the data.
294// The exact table necessary is given in the OpenDML spec.""
295pub fn fill_default_mjpeg_tables(scan: &ScanInfo,
296 dc_huffman_tables: &mut[Option<HuffmanTable>],
297 ac_huffman_tables: &mut[Option<HuffmanTable>]) {
298 // Section K.3.3
299
300 if dc_huffman_tables[0].is_none() && scan.dc_table_indices.iter().any(|&i| i == 0) {
301 // Table K.3
302 dc_huffman_tables[0] = Some(HuffmanTable::new(
303 &[0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
304 &[0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B], HuffmanTableClass::DC).unwrap());
305 }
306 if dc_huffman_tables[1].is_none() && scan.dc_table_indices.iter().any(|&i| i == 1) {
307 // Table K.4
308 dc_huffman_tables[1] = Some(HuffmanTable::new(
309 &[0x00, 0x03, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00],
310 &[0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B], HuffmanTableClass::DC).unwrap());
311 }
312 if ac_huffman_tables[0].is_none() && scan.ac_table_indices.iter().any(|&i| i == 0) {
313 // Table K.5
314 ac_huffman_tables[0] = Some(HuffmanTable::new(
315 &[0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03, 0x05, 0x05, 0x04, 0x04, 0x00, 0x00, 0x01, 0x7D],
316 &[0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
317 0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xA1, 0x08, 0x23, 0x42, 0xB1, 0xC1, 0x15, 0x52, 0xD1, 0xF0,
318 0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28,
319 0x29, 0x2A, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
320 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69,
321 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
322 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7,
323 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5,
324 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA, 0xE1, 0xE2,
325 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
326 0xF9, 0xFA
327 ], HuffmanTableClass::AC).unwrap());
328 }
329 if ac_huffman_tables[1].is_none() && scan.ac_table_indices.iter().any(|&i| i == 1) {
330 // Table K.6
331 ac_huffman_tables[1] = Some(HuffmanTable::new(
332 &[0x00, 0x02, 0x01, 0x02, 0x04, 0x04, 0x03, 0x04, 0x07, 0x05, 0x04, 0x04, 0x00, 0x01, 0x02, 0x77],
333 &[0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21, 0x31, 0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71,
334 0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1, 0x09, 0x23, 0x33, 0x52, 0xF0,
335 0x15, 0x62, 0x72, 0xD1, 0x0A, 0x16, 0x24, 0x34, 0xE1, 0x25, 0xF1, 0x17, 0x18, 0x19, 0x1A, 0x26,
336 0x27, 0x28, 0x29, 0x2A, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48,
337 0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68,
338 0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
339 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5,
340 0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3,
341 0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA,
342 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
343 0xF9, 0xFA
344 ], HuffmanTableClass::AC).unwrap());
345 }
346}
347