zune_jpeg/
huffman.rs

1/*
2 * Copyright (c) 2023.
3 *
4 * This software is free software;
5 *
6 * You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
7 */
8
9//! This file contains a single struct `HuffmanTable` that
10//! stores Huffman tables needed during `BitStream` decoding.
11#![allow(clippy::similar_names, clippy::module_name_repetitions)]
12
13use alloc::string::ToString;
14
15use crate::errors::DecodeErrors;
16
17/// Determines how many bits of lookahead we have for our bitstream decoder.
18
19pub const HUFF_LOOKAHEAD: u8 = 9;
20
21/// A struct which contains necessary tables for decoding a JPEG
22/// huffman encoded bitstream
23
24pub struct HuffmanTable {
25    // element `[0]` of each array is unused
26    /// largest code of length k
27    pub(crate) maxcode: [i32; 18],
28    /// offset for codes of length k
29    /// Answers the question, where do code-lengths of length k end
30    /// Element 0 is unused
31    pub(crate) offset:  [i32; 18],
32    /// lookup table for fast decoding
33    ///
34    /// top  bits above HUFF_LOOKAHEAD contain the code length.
35    ///
36    /// Lower (8) bits contain the symbol in order of increasing code length.
37    pub(crate) lookup:  [i32; 1 << HUFF_LOOKAHEAD],
38
39    /// A table which can be used to decode small AC coefficients and
40    /// do an equivalent of receive_extend
41    pub(crate) ac_lookup: Option<[i16; 1 << HUFF_LOOKAHEAD]>,
42
43    /// Directly represent contents of a JPEG DHT marker
44    ///
45    /// \# number of symbols with codes of length `k` bits
46    // bits[0] is unused
47    /// Symbols in order of increasing code length
48    pub(crate) values: [u8; 256]
49}
50
51impl HuffmanTable {
52    pub fn new(
53        codes: &[u8; 17], values: [u8; 256], is_dc: bool, is_progressive: bool
54    ) -> Result<HuffmanTable, DecodeErrors> {
55        let too_long_code = (i32::from(HUFF_LOOKAHEAD) + 1) << HUFF_LOOKAHEAD;
56        let mut p = HuffmanTable {
57            maxcode: [0; 18],
58            offset: [0; 18],
59            lookup: [too_long_code; 1 << HUFF_LOOKAHEAD],
60            values,
61            ac_lookup: None
62        };
63
64        p.make_derived_table(is_dc, is_progressive, codes)?;
65
66        Ok(p)
67    }
68
69    /// Create a new huffman tables with values that aren't fixed
70    /// used by fill_mjpeg_tables
71    pub fn new_unfilled(
72        codes: &[u8; 17], values: &[u8], is_dc: bool, is_progressive: bool
73    ) -> Result<HuffmanTable, DecodeErrors> {
74        let mut buf = [0; 256];
75        buf[..values.len()].copy_from_slice(values);
76        HuffmanTable::new(codes, buf, is_dc, is_progressive)
77    }
78
79    /// Compute derived values for a Huffman table
80    ///
81    /// This routine performs some validation checks on the table
82    #[allow(
83        clippy::cast_possible_truncation,
84        clippy::cast_possible_wrap,
85        clippy::cast_sign_loss,
86        clippy::too_many_lines,
87        clippy::needless_range_loop
88    )]
89    fn make_derived_table(
90        &mut self, is_dc: bool, _is_progressive: bool, bits: &[u8; 17]
91    ) -> Result<(), DecodeErrors> {
92        // build a list of code size
93        let mut huff_size = [0; 257];
94        // Huffman code lengths
95        let mut huff_code: [u32; 257] = [0; 257];
96        // figure C.1 make table of Huffman code length for each symbol
97        let mut p = 0;
98
99        for l in 1..=16 {
100            let mut i = i32::from(bits[l]);
101            // table overrun is checked before ,so we dont need to check
102            while i != 0 {
103                huff_size[p] = l as u8;
104                p += 1;
105                i -= 1;
106            }
107        }
108
109        huff_size[p] = 0;
110
111        let num_symbols = p;
112        // Generate the codes themselves
113        // We also validate that the counts represent a legal Huffman code tree
114        let mut code = 0;
115        let mut si = i32::from(huff_size[0]);
116
117        p = 0;
118
119        while huff_size[p] != 0 {
120            while i32::from(huff_size[p]) == si {
121                huff_code[p] = code;
122                code += 1;
123                p += 1;
124            }
125            // maximum code of length si, pre-shifted by 16-k bits
126            self.maxcode[si as usize] = (code << (16 - si)) as i32;
127            // code is now 1 more than the last code used for code-length si; but
128            // it must still fit in si bits, since no code is allowed to be all ones.
129            if (code as i32) >= (1 << si) {
130                return Err(DecodeErrors::HuffmanDecode("Bad Huffman Table".to_string()));
131            }
132
133            code <<= 1;
134            si += 1;
135        }
136
137        // Figure F.15 generate decoding tables for bit-sequential decoding
138        p = 0;
139
140        for l in 0..=16 {
141            if bits[l] == 0 {
142                // -1 if no codes of this length
143                self.maxcode[l] = -1;
144            } else {
145                // offset[l]=codes[index of 1st symbol of code length l
146                // minus minimum code of length l]
147                self.offset[l] = (p as i32) - (huff_code[p]) as i32;
148                p += usize::from(bits[l]);
149            }
150        }
151
152        self.offset[17] = 0;
153        // we ensure that decode terminates
154        self.maxcode[17] = 0x000F_FFFF;
155
156        /*
157         * Compute lookahead tables to speed up decoding.
158         * First we set all the table entries to 0(left justified), indicating "too long";
159         * (Note too long was set during initialization)
160         * then we iterate through the Huffman codes that are short enough and
161         * fill in all the entries that correspond to bit sequences starting
162         * with that code.
163         */
164
165        p = 0;
166
167        for l in 1..=HUFF_LOOKAHEAD {
168            for _ in 1..=i32::from(bits[usize::from(l)]) {
169                // l -> Current code length,
170                // p => Its index in self.code and self.values
171                // Generate left justified code followed by all possible bit sequences
172                let mut look_bits = (huff_code[p] as usize) << (HUFF_LOOKAHEAD - l);
173
174                for _ in 0..1 << (HUFF_LOOKAHEAD - l) {
175                    self.lookup[look_bits] =
176                        (i32::from(l) << HUFF_LOOKAHEAD) | i32::from(self.values[p]);
177                    look_bits += 1;
178                }
179
180                p += 1;
181            }
182        }
183        // build an ac table that does an equivalent of decode and receive_extend
184        if !is_dc {
185            let mut fast = [255; 1 << HUFF_LOOKAHEAD];
186            // Iterate over number of symbols
187            for i in 0..num_symbols {
188                // get code size for an item
189                let s = huff_size[i];
190
191                if s <= HUFF_LOOKAHEAD {
192                    // if it's lower than what we need for our lookup table create the table
193                    let c = (huff_code[i] << (HUFF_LOOKAHEAD - s)) as usize;
194                    let m = (1 << (HUFF_LOOKAHEAD - s)) as usize;
195
196                    for j in 0..m {
197                        fast[c + j] = i as i16;
198                    }
199                }
200            }
201
202            // build a table that decodes both magnitude and value of small ACs in
203            // one go.
204            let mut fast_ac = [0; 1 << HUFF_LOOKAHEAD];
205
206            for i in 0..(1 << HUFF_LOOKAHEAD) {
207                let fast_v = fast[i];
208
209                if fast_v < 255 {
210                    // get symbol value from AC table
211                    let rs = self.values[fast_v as usize];
212                    // shift by 4 to get run length
213                    let run = i16::from((rs >> 4) & 15);
214                    // get magnitude bits stored at the lower 3 bits
215                    let mag_bits = i16::from(rs & 15);
216                    // length of the bit we've read
217                    let len = i16::from(huff_size[fast_v as usize]);
218
219                    if mag_bits != 0 && (len + mag_bits) <= i16::from(HUFF_LOOKAHEAD) {
220                        // magnitude code followed by receive_extend code
221                        let mut k = (((i as i16) << len) & ((1 << HUFF_LOOKAHEAD) - 1))
222                            >> (i16::from(HUFF_LOOKAHEAD) - mag_bits);
223                        let m = 1 << (mag_bits - 1);
224
225                        if k < m {
226                            k += (!0_i16 << mag_bits) + 1;
227                        };
228
229                        // if result is small enough fit into fast ac table
230                        if (-128..=127).contains(&k) {
231                            fast_ac[i] = (k << 8) + (run << 4) + (len + mag_bits);
232                        }
233                    }
234                }
235            }
236            self.ac_lookup = Some(fast_ac);
237        }
238
239        // Validate symbols as being reasonable
240        // For AC tables, we make no check, but accept all byte values 0..255
241        // For DC tables, we require symbols to be in range 0..15
242        if is_dc {
243            for i in 0..num_symbols {
244                let sym = self.values[i];
245
246                if sym > 15 {
247                    return Err(DecodeErrors::HuffmanDecode("Bad Huffman Table".to_string()));
248                }
249            }
250        }
251
252        Ok(())
253    }
254}