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