flashraw_core/
huffman.rs

1//! Canonical Huffman decoding (JPEG DHT style) with lookup-table acceleration.
2//!
3//! Tables are described the JPEG way: `counts[i]` codes of length `i + 1` followed
4//! by the symbol values in code order. Most RAW codecs (lossless JPEG, Nikon,
5//! Pentax, Sony ARW1, Kodak, Hasselblad, Leaf, Samsung...) use this representation
6//! with symbols that are the bit-length of a difference value.
7//!
8//! The decoder uses a single `LUT_BITS`-bit lookup. Each entry either
9//! - holds the *fully decoded* difference (code + diff bits fit in the LUT), or
10//! - holds the code length and symbol (diff bits are then read separately), or
11//! - flags a long code that falls back to the canonical slow path.
12
13use crate::bits::{BitRead, extend_diff};
14
15/// Sign-extend a JPEG difference without branches (`len` in 0..=15).
16#[inline(always)]
17fn extend_diff_branchless(v: u32, len: u32) -> i32 {
18    // For len == 0, v == 0 and the shifted test bit is 0 → neg mask applies to
19    // (1 - 1) == 0, giving 0 without a branch.
20    let v = v as i32;
21    let top = (v >> (len.max(1) - 1)) & 1;
22    // All ones when the top bit is clear (negative difference), else zero.
23    let neg = top - 1;
24    v + (neg & (1 - (1i32 << len)))
25}
26
27#[cfg(test)]
28#[test]
29fn branchless_extension_matches_reference() {
30    for len in 0..16u32 {
31        for v in 0..(1u32 << len) {
32            assert_eq!(extend_diff_branchless(v, len), extend_diff(v, len), "{v} {len}");
33        }
34    }
35}
36use crate::{Error, Result};
37
38const LUT_BITS: u32 = 12;
39const LUT_SIZE: usize = 1 << LUT_BITS;
40
41// Entry layout (u32):
42//   bits 0..5   : bits consumed: the code length, or code + diff length when FULL
43//   bit  5      : FULL flag (bits 16..32 hold the decoded difference)
44//   bit  6      : VALID flag (code fits in LUT_BITS; otherwise use the canonical slow path)
45//   bits 8..16  : symbol (for difference tables: the number of diff bits that follow)
46//   bits 16..32 : decoded difference (i16) when FULL
47//
48// Two decoding strategies, chosen per table from its code lengths:
49// - "full" tables (most of the probability mass decodes in one lookup, e.g. short codes
50//   with small differences): FULL entries carry the finished difference; the FULL
51//   branch is then well predicted and one load decodes a sample.
52// - otherwise entries store only code length + symbol and the difference is extracted
53//   arithmetically (peek code + diff bits at once, mask, branch-free sign extension),
54//   avoiding mispredictions when the FULL case is a coin flip.
55const FLAG_FULL: u32 = 1 << 5;
56const FLAG_VALID: u32 = 1 << 6;
57/// Probability mass of single-lookup samples above which a table uses FULL entries.
58const FULL_MODE_THRESHOLD: f64 = 0.0;
59
60/// A canonical Huffman table with fast lookup.
61#[derive(Clone)]
62pub struct HuffTable {
63    lut: Box<[u32; LUT_SIZE]>,
64    /// For the slow path: max code (exclusive upper bound, left-aligned per length), per length 1..=16.
65    maxcode: [i32; 18],
66    valptr: [i32; 17],
67    mincode: [u32; 17],
68    values: Vec<u8>,
69    /// Lossless JPEG/DNG quirk: a length-16 symbol means diff = -32768 with no extra bits.
70    pub dng_bug16: bool,
71    /// Whether diff bits follow the symbol (false: the symbol value itself is the result).
72    pub has_diffs: bool,
73    /// Use precomputed differences (see the entry layout notes).
74    full_mode: bool,
75    max_len: u32,
76}
77
78impl std::fmt::Debug for HuffTable {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        f.debug_struct("HuffTable").field("values", &self.values).field("max_len", &self.max_len).finish()
81    }
82}
83
84impl HuffTable {
85    /// Build from JPEG-style `counts` (codes per length 1..=16) and `values`.
86    pub fn from_counts(counts: &[u8], values: &[u8], has_diffs: bool) -> Result<Self> {
87        if counts.len() > 16 {
88            return Err(Error::corrupt("Huffman table with code lengths > 16"));
89        }
90        let total: usize = counts.iter().map(|&c| c as usize).sum();
91        if total == 0 || total > values.len() || total > 256 {
92            return Err(Error::corrupt(format!("invalid Huffman table ({total} codes, {} values)", values.len())));
93        }
94        let values = values[..total].to_vec();
95        // Generate canonical codes.
96        let mut maxcode = [-1i32; 18];
97        let mut valptr = [0i32; 17];
98        let mut mincode = [0u32; 17];
99        let mut code: u32 = 0;
100        let mut k = 0usize;
101        let mut max_len = 0;
102        let mut codes: Vec<(u32, u32, u8)> = Vec::with_capacity(total); // (code, len, symbol)
103        for len in 1..=16usize {
104            let n = counts.get(len - 1).copied().unwrap_or(0) as usize;
105            if n > 0 {
106                valptr[len] = k as i32;
107                mincode[len] = code;
108                for _ in 0..n {
109                    if code >= (1u32 << len) {
110                        return Err(Error::corrupt("over-subscribed Huffman table"));
111                    }
112                    codes.push((code, len as u32, values[k]));
113                    code += 1;
114                    k += 1;
115                }
116                maxcode[len] = code as i32 - 1;
117                max_len = len as u32;
118            }
119            code <<= 1;
120        }
121        maxcode[17] = i32::MAX;
122        // Expected share of samples whose code + difference fit one lookup (a code of
123        // length L occurs with probability ~2^-L in a canonical, near-optimal code).
124        let full_mass: f64 = codes
125            .iter()
126            .filter(|&&(_, len, sym)| len <= LUT_BITS && (sym == 0 || (sym < 16 && len + sym as u32 <= LUT_BITS)))
127            .map(|&(_, len, _)| 0.5f64.powi(len as i32))
128            .sum();
129        let full_mode = has_diffs && full_mass >= FULL_MODE_THRESHOLD;
130        let mut lut = Box::new([0u32; LUT_SIZE]);
131        for &(code, len, sym) in &codes {
132            if len > LUT_BITS {
133                continue;
134            }
135            let shift = LUT_BITS - len;
136            let base = (code << shift) as usize;
137            let plain = len | FLAG_VALID | ((sym as u32) << 8);
138            if !full_mode {
139                lut[base..base + (1usize << shift)].fill(plain);
140                continue;
141            }
142            for fill in 0..(1usize << shift) {
143                let idx = base | fill;
144                let dlen = sym as u32;
145                lut[idx] = if dlen == 0 {
146                    plain | FLAG_FULL
147                } else if dlen < 16 && len + dlen <= LUT_BITS {
148                    let dbits = ((idx as u32) >> (LUT_BITS - len - dlen)) & ((1 << dlen) - 1);
149                    let diff = extend_diff(dbits, dlen);
150                    (len + dlen) | FLAG_VALID | FLAG_FULL | ((sym as u32) << 8) | ((diff as i16 as u16 as u32) << 16)
151                } else {
152                    plain
153                };
154            }
155        }
156        Ok(HuffTable { lut, maxcode, valptr, mincode, values, dng_bug16: false, has_diffs, full_mode, max_len })
157    }
158
159    /// Build from a list of `(code_length, symbol)` pairs in canonical order
160    /// (helper for vendors that specify tables as code lengths).
161    pub fn from_lengths(lengths: &[(u8, u8)], has_diffs: bool) -> Result<Self> {
162        let mut counts = [0u8; 16];
163        let mut sorted: Vec<(u8, u8)> = lengths.to_vec();
164        sorted.sort_by_key(|&(l, _)| l);
165        let mut values = Vec::with_capacity(sorted.len());
166        for &(l, s) in &sorted {
167            if l == 0 || l > 16 {
168                return Err(Error::corrupt("invalid Huffman code length"));
169            }
170            counts[l as usize - 1] += 1;
171            values.push(s);
172        }
173        Self::from_counts(&counts, &values, has_diffs)
174    }
175
176    /// Maximum code length in this table.
177    pub fn max_len(&self) -> u32 {
178        self.max_len
179    }
180
181    #[cold]
182    fn slow_symbol<B: BitRead>(&self, bits: &mut B) -> Result<(u32, u32)> {
183        // Canonical decode bit by bit beyond LUT_BITS.
184        bits.fill();
185        let peek = bits.peek_no_fill(16);
186        for len in 1..=16u32 {
187            let code = (peek >> (16 - len)) as i32;
188            if self.maxcode[len as usize] >= 0
189                && code <= self.maxcode[len as usize]
190                && code as u32 >= self.mincode[len as usize]
191            {
192                let idx = self.valptr[len as usize] + code - self.mincode[len as usize] as i32;
193                let sym = *self.values.get(idx as usize).ok_or_else(|| Error::corrupt("invalid Huffman index"))?;
194                return Ok((sym as u32, len));
195            }
196        }
197        Err(Error::corrupt("invalid Huffman code"))
198    }
199
200    /// Decode one symbol (no diff handling).
201    #[inline(always)]
202    pub fn decode_symbol<B: BitRead>(&self, bits: &mut B) -> Result<u32> {
203        bits.fill();
204        let e = self.lut[bits.peek_no_fill(LUT_BITS) as usize];
205        if e & FLAG_VALID != 0 {
206            let sym = (e >> 8) & 0xFF;
207            // FULL entries count code + diff bits; skip only the code here.
208            let code_len = if e & FLAG_FULL != 0 { (e & 0x1F) - sym } else { e & 0x1F };
209            bits.skip_no_fill(code_len);
210            return Ok(sym);
211        }
212        let (sym, len) = self.slow_symbol(bits)?;
213        bits.skip_no_fill(len);
214        Ok(sym)
215    }
216
217    /// Decode a symbol giving a bit length, then read and sign-extend that many diff bits.
218    #[inline(always)]
219    pub fn decode_diff<B: BitRead>(&self, bits: &mut B) -> Result<i32> {
220        bits.fill();
221        let e = self.lut[bits.peek_no_fill(LUT_BITS) as usize];
222        if self.full_mode && e & FLAG_FULL != 0 {
223            bits.skip_no_fill(e & 0x1F);
224            return Ok((e >> 16) as u16 as i16 as i32);
225        }
226        let dlen = (e >> 8) & 0xFF;
227        if e & FLAG_VALID != 0 && dlen < 16 {
228            // After fill() at least 32 bits are buffered; code (<= LUT_BITS) + diff (<= 15)
229            // always fits, so code and diff are consumed together, without branches.
230            let total = (e & 0x1F) + dlen;
231            let v = bits.peek_no_fill(total) & ((1u32 << dlen) - 1);
232            bits.skip_no_fill(total);
233            return Ok(extend_diff_branchless(v, dlen));
234        }
235        let (sym, len) = if e & FLAG_VALID != 0 { (dlen, e & 0x1F) } else { self.slow_symbol(bits)? };
236        bits.skip_no_fill(len);
237        Ok(self.read_diff(bits, sym))
238    }
239
240    /// Decode one difference assuming at least 27 bits are buffered (the caller refilled,
241    /// e.g. with [`BitRead::fill_to`]):
242    /// table hits need no refill and no error path. Rare cases (codes longer than the
243    /// lookup table, 16-bit differences, invalid codes) fall back to the general decoder;
244    /// a corrupt code sets `*err` and yields 0 so hot loops avoid per-sample `Result`s.
245    #[inline(always)]
246    pub fn decode_diff_buffered<B: BitRead>(&self, bits: &mut B, err: &mut bool) -> i32 {
247        let e = self.lut[bits.peek_no_fill(LUT_BITS) as usize];
248        if e & FLAG_FULL != 0 {
249            bits.skip_no_fill(e & 0x1F);
250            return (e >> 16) as u16 as i16 as i32;
251        }
252        let dlen = (e >> 8) & 0xFF;
253        if e & FLAG_VALID != 0 && dlen < 16 {
254            let total = (e & 0x1F) + dlen;
255            let v = bits.peek_no_fill(total) & ((1u32 << dlen) - 1);
256            bits.skip_no_fill(total);
257            return extend_diff_branchless(v, dlen);
258        }
259        self.decode_diff_cold(bits, err)
260    }
261
262    #[cold]
263    #[inline(never)]
264    fn decode_diff_cold<B: BitRead>(&self, bits: &mut B, err: &mut bool) -> i32 {
265        match self.decode_diff(bits) {
266            Ok(d) => d,
267            Err(_) => {
268                *err = true;
269                0
270            }
271        }
272    }
273
274    /// Read the diff for a given length symbol (handles the 16-bit special case).
275    #[inline(always)]
276    pub fn read_diff<B: BitRead>(&self, bits: &mut B, len: u32) -> i32 {
277        match len {
278            0 => 0,
279            16 if self.dng_bug16 => -32768,
280            16 => {
281                // Plain 16-bit diff: no sign-extension ambiguity beyond 16 bits.
282                let v = bits.get(16);
283                extend_diff(v, 16)
284            }
285            l if l > 16 => {
286                // Out-of-spec lengths (seen in corrupt data); read and clamp safely.
287                let v = bits.get(16);
288                extend_diff(v, 16)
289            }
290            l => {
291                let v = bits.get(l);
292                extend_diff(v, l)
293            }
294        }
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use crate::bits::BitPumpMsb;
302
303    struct Writer {
304        out: Vec<u8>,
305        acc: u64,
306        n: u32,
307    }
308    impl Writer {
309        fn put(&mut self, v: u32, len: u32) {
310            for i in (0..len).rev() {
311                self.acc = (self.acc << 1) | ((v >> i) & 1) as u64;
312                self.n += 1;
313                if self.n == 8 {
314                    self.out.push(self.acc as u8);
315                    self.acc = 0;
316                    self.n = 0;
317                }
318            }
319        }
320        fn finish(mut self) -> Vec<u8> {
321            while self.n != 0 {
322                self.put(0, 1);
323            }
324            self.out.extend_from_slice(&[0; 8]);
325            self.out
326        }
327    }
328
329    #[test]
330    fn roundtrip_diffs_including_long_codes() {
331        // Table with codes up to length 16 to exercise the slow path.
332        let counts = [0u8, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1];
333        let values: Vec<u8> = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
334        let t = HuffTable::from_counts(&counts, &values, true).unwrap();
335        // Rebuild canonical codes for encoding.
336        let mut codes = vec![];
337        let mut code = 0u32;
338        let mut k = 0;
339        for len in 1..=16u32 {
340            for _ in 0..counts[len as usize - 1] {
341                codes.push((values[k], code, len));
342                code += 1;
343                k += 1;
344            }
345            code <<= 1;
346        }
347        let diffs: Vec<i32> =
348            (0..500).map(|i: i32| ((i * 7919) % 20001) - 10000).map(|d| d / (1 + (d.abs() % 7))).collect();
349        let mut w = Writer { out: vec![], acc: 0, n: 0 };
350        for &d in &diffs {
351            let len = if d == 0 { 0 } else { 32 - d.unsigned_abs().leading_zeros() };
352            let (_, c, l) = codes.iter().find(|(s, _, _)| *s as u32 == len).unwrap();
353            w.put(*c, *l);
354            if len > 0 {
355                let v = if d < 0 { (d - 1) as u32 & ((1 << len) - 1) } else { d as u32 };
356                w.put(v, len);
357            }
358        }
359        let data = w.finish();
360        let mut p = BitPumpMsb::new(&data);
361        for &d in &diffs {
362            assert_eq!(t.decode_diff(&mut p).unwrap(), d);
363        }
364    }
365
366    #[test]
367    fn rejects_bad_tables() {
368        assert!(HuffTable::from_counts(&[3], &[0, 1, 2], true).is_err());
369        assert!(HuffTable::from_counts(&[0; 16], &[], true).is_err());
370    }
371}
flashraw_core/huffman.rs

flashraw_core/
huffman.rs
