jpeg_decoder/
huffman.rs

Help
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 = 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}
jpeg_decoder/huffman.rs

jpeg_decoder/
huffman.rs
