Skip to main content

flashraw_codecs/
nikon_huffman.rs

1//! Nikon NEF Huffman compression, plus the predictive Huffman engine
2//! it shares with Pentax PEF and Samsung SRW (compression 32772).
3//!
4//! # Bitstream model
5//!
6//! All three formats code a CFA image row by row as a single MSB-first bit stream
7//! (no JPEG byte stuffing). Every sample is a Huffman symbol followed by a
8//! difference; the predictor is the previous same-colour sample of the row
9//! (`hpred[col & 1]`), and the first two samples of a row predict from the first two
10//! samples two rows above (`vpred[row & 1][col]`).
11//!
12//! Symbols are stored in byte form: low nibble = total difference length
13//! `len`, high nibble = `shl`, the number of low-order difference bits that are *not*
14//! transmitted (Nikon's lossy "after split" tables). For `shl == 0` the difference is
15//! the usual JPEG sign-extended value.
16//!
17//! # Truncated streams
18//!
19//! Decoding reads past the end of the strip into the rest of the file and, at end
20//! of file, keeps decoding with a zero-padded bit buffer
21//! until the first read that crosses EOF; every later read returns zero (so the
22//! predictors freeze). [`ExactBits`] implements these semantics so truncated files
23//! decode deterministically.
24//!
25//! # Performance
26//!
27//! The hot loop uses one unaligned 64-bit big-endian load per sample: a
28//! [`FAST_BITS`]-bit lookup either yields the fully decoded difference (code plus
29//! difference bits fit in the index) or the code length and symbol, whose difference
30//! bits are taken from the same load. Rows that are provably far from the end of the
31//! data use this path; the last rows go through the exact slow path.
32
33use flashraw_core::bytes::ByteOrder;
34use flashraw_core::{Error, Result};
35
36/// Index width of the combined code+difference lookup table.
37pub const FAST_BITS: u32 = 13;
38const FAST_SIZE: usize = 1 << FAST_BITS;
39const FLAG_FULL: u32 = 1 << 5;
40/// Fast-table entry whose code is longer than `FAST_BITS` (resolve via the direct table).
41const SLOW_MARK: u32 = 0x8000_0000;
42
43/// A Huffman table in "direct lookup" form: `peek_bits` bits index a table of
44/// `code_len << 8 | symbol` entries (zero for unused codes, which consume no bits and
45/// decode to symbol 0, which is how incomplete tables behave).
46#[derive(Clone)]
47pub struct DirectHuff {
48    peek_bits: u32,
49    direct: Vec<u16>,
50    /// Combined table indexed by the next `FAST_BITS` bits. Entry layout:
51    /// bits 0..5 bits consumed; bit 5 FULL (diff in bits 16..32); bits 8..16 symbol.
52    fast: Box<[u32; FAST_SIZE]>,
53}
54
55impl std::fmt::Debug for DirectHuff {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        f.debug_struct("DirectHuff").field("peek_bits", &self.peek_bits).finish()
58    }
59}
60
61/// Difference value for a symbol (`len | shl << 4`) given the `len - shl` transmitted bits.
62#[inline(always)]
63pub fn symbol_diff(bits: u32, sym: u32) -> i32 {
64    let len = sym & 15;
65    let shl = sym >> 4;
66    let mut diff = ((((bits as i32) << 1) + 1) << shl) >> 1;
67    if len > 0 && diff & (1 << (len - 1)) == 0 {
68        diff -= (1 << len) - (shl == 0) as i32;
69    }
70    diff
71}
72
73impl DirectHuff {
74    /// Build from a table source: 16 code-length counts followed by the
75    /// symbols in code order.
76    pub fn from_counts(source: &[u8]) -> Result<Self> {
77        if source.len() < 16 {
78            return Err(Error::corrupt("Huffman source shorter than 16 counts"));
79        }
80        let counts = &source[..16];
81        let max = (1..=16).rev().find(|&l| counts[l - 1] != 0).unwrap_or(0) as u32;
82        if max == 0 {
83            return Err(Error::corrupt("empty Huffman table"));
84        }
85        let total: usize = counts.iter().map(|&c| c as usize).sum();
86        let syms = source.get(16..16 + total).ok_or_else(|| Error::corrupt("missing Huffman table symbols"))?;
87        let mut direct = vec![0u16; 1 << max];
88        let mut h = 0usize;
89        let mut k = 0usize;
90        for len in 1..=max {
91            for _ in 0..counts[len as usize - 1] {
92                let sym = syms[k];
93                k += 1;
94                for _ in 0..(1usize << (max - len)) {
95                    if h < direct.len() {
96                        direct[h] = ((len as u16) << 8) | sym as u16;
97                        h += 1;
98                    }
99                }
100            }
101        }
102        Ok(Self::from_direct(max, direct))
103    }
104
105    /// Build from a direct table of `1 << peek_bits` entries (`code_len << 8 | symbol`).
106    pub fn from_direct(peek_bits: u32, direct: Vec<u16>) -> Self {
107        debug_assert_eq!(direct.len(), 1 << peek_bits);
108        let mut fast: Box<[u32; FAST_SIZE]> = vec![0u32; FAST_SIZE].into_boxed_slice().try_into().expect("FAST_SIZE");
109        for (idx, e) in fast.iter_mut().enumerate() {
110            // Code bits are the top `peek_bits` of the FAST_BITS index. When the table is
111            // wider than the fast index, only codes short enough to be determined by
112            // FAST_BITS bits are resolved here (canonical tables fill aligned blocks);
113            // the rest are marked for the slow path.
114            let entry = if peek_bits <= FAST_BITS {
115                direct[idx >> (FAST_BITS - peek_bits)]
116            } else {
117                let d = direct[idx << (peek_bits - FAST_BITS)];
118                let clen = (d >> 8) as u32;
119                if clen == 0 || clen > FAST_BITS {
120                    *e = SLOW_MARK;
121                    continue;
122                }
123                d
124            };
125            let clen = (entry >> 8) as u32;
126            let sym = (entry & 0xFF) as u32;
127            let len = sym & 15;
128            let shl = sym >> 4;
129            let nbits = len.saturating_sub(shl);
130            if clen + nbits <= FAST_BITS {
131                let b = if nbits == 0 { 0 } else { ((idx as u32) >> (FAST_BITS - clen - nbits)) & ((1 << nbits) - 1) };
132                let d = symbol_diff(b, sym);
133                *e = (clen + nbits) | FLAG_FULL | (sym << 8) | ((d as i16 as u16 as u32) << 16);
134            } else {
135                *e = clen | (sym << 8);
136            }
137        }
138        DirectHuff { peek_bits, direct, fast }
139    }
140
141    pub fn peek_bits(&self) -> u32 {
142        self.peek_bits
143    }
144
145    /// Direct-table entry for the next `peek_bits` bits.
146    #[inline(always)]
147    fn direct(&self, idx: u32) -> u16 {
148        self.direct[idx as usize]
149    }
150}
151
152/// MSB-first bit reader with exact end-of-data semantics (see module docs).
153#[derive(Clone)]
154pub struct ExactBits<'a> {
155    data: &'a [u8],
156    /// Bits consumed so far.
157    pos: u64,
158    /// Set after the first read that crossed the end of the data.
159    dead: bool,
160    /// First output row during which the stream ran past its data ([`decode_predictive`]).
161    dead_row: Option<usize>,
162}
163
164impl<'a> ExactBits<'a> {
165    pub fn new(data: &'a [u8]) -> Self {
166        ExactBits { data, pos: 0, dead: false, dead_row: None }
167    }
168
169    /// Bits consumed so far.
170    pub fn position(&self) -> u64 {
171        self.pos
172    }
173
174    /// Whether the stream ran past the end of its data.
175    pub fn is_dead(&self) -> bool {
176        self.dead
177    }
178
179    /// The first row (of a [`decode_predictive`] call) that needed bits past the end of the
180    /// data, i.e. how much of the image is real: rows from here on are frozen predictors.
181    pub fn dead_row(&self) -> Option<usize> {
182        self.dead_row
183    }
184
185    #[inline(always)]
186    fn total_bits(&self) -> u64 {
187        self.data.len() as u64 * 8
188    }
189
190    /// The 57+ bits starting at the current position, left-aligned, zero padded past the end.
191    #[inline(always)]
192    fn window(&self) -> u64 {
193        let byte = (self.pos >> 3) as usize;
194        let w = match self.data.get(byte..byte + 8) {
195            Some(s) => u64::from_be_bytes(s.try_into().unwrap()),
196            None => {
197                let mut b = [0u8; 8];
198                if let Some(s) = self.data.get(byte..) {
199                    b[..s.len()].copy_from_slice(s);
200                }
201                u64::from_be_bytes(b)
202            }
203        };
204        w << (self.pos & 7)
205    }
206
207    /// Consume `n` bits after peeking (marks the stream dead if this crosses the end).
208    #[inline(always)]
209    fn consume(&mut self, n: u32) {
210        self.pos += n as u64;
211        if self.pos > self.total_bits() {
212            self.dead = true;
213        }
214    }
215
216    /// Read `n <= 25` bits.
217    #[inline]
218    pub fn get(&mut self, n: u32) -> u32 {
219        if n == 0 || self.dead {
220            return 0;
221        }
222        let v = (self.window() >> (64 - n)) as u32;
223        self.consume(n);
224        v
225    }
226
227    /// Decode one symbol.
228    #[inline]
229    pub fn huff(&mut self, t: &DirectHuff) -> u32 {
230        if self.dead {
231            return 0;
232        }
233        let idx = (self.window() >> (64 - t.peek_bits)) as u32;
234        let e = t.direct(idx);
235        self.consume((e >> 8) as u32);
236        (e & 0xFF) as u32
237    }
238
239    /// Decode one symbol + difference (exact slow path).
240    #[inline]
241    pub fn diff(&mut self, t: &DirectHuff) -> i32 {
242        let sym = self.huff(t);
243        let len = sym & 15;
244        let shl = sym >> 4;
245        let b = self.get(len.saturating_sub(shl));
246        symbol_diff(b, sym)
247    }
248}
249
250/// Parameters of a predictive Huffman image.
251pub struct PredictiveParams<'t> {
252    pub width: usize,
253    /// Rows to decode (rows beyond are left untouched).
254    pub rows: usize,
255    /// Tables and the first row each applies to (ascending; first must start at 0).
256    pub tables: &'t [(usize, &'t DirectHuff)],
257    /// Initial vertical predictors (`vpred[row & 1][col]` for the first two columns).
258    pub vpred: [[u16; 2]; 2],
259}
260
261/// Decode a predictive Huffman image into `out` (rows of `stride` samples), mapping each
262/// reconstructed predictor value through `map` (identity, or a linearisation curve).
263///
264/// Returns the reader so callers can inspect truncation.
265pub fn decode_predictive<'a, M>(
266    src: &'a [u8],
267    p: &PredictiveParams,
268    out: &mut [u16],
269    stride: usize,
270    map: M,
271) -> Result<ExactBits<'a>>
272where
273    M: Fn(u16) -> u16,
274{
275    let PredictiveParams { width, rows, tables, vpred } = *p;
276    if width < 2 || stride < width {
277        return Err(Error::InvalidArgument("predictive Huffman: invalid width".into()));
278    }
279    if out.len() < stride * rows.saturating_sub(1) + width && rows > 0 {
280        return Err(Error::InvalidArgument("predictive Huffman: output too small".into()));
281    }
282    let first = tables.first().ok_or_else(|| Error::InvalidArgument("no Huffman table".into()))?;
283    if first.0 != 0 {
284        return Err(Error::InvalidArgument("first Huffman table must start at row 0".into()));
285    }
286    let mut vpred = vpred;
287    let mut bits = ExactBits::new(src);
288    let mut ti = 0usize;
289    // Worst case bits per sample is 31 (16-bit code + 15 difference bits); the fast
290    // reader also looks up to 16 bytes ahead of its position.
291    let row_budget = width as u64 * 32 + 192;
292    for row in 0..rows {
293        while ti + 1 < tables.len() && tables[ti + 1].0 <= row {
294            ti += 1;
295        }
296        let t = tables[ti].1;
297        let line = &mut out[row * stride..row * stride + width];
298        let vp = &mut vpred[row & 1];
299        let fast = !bits.dead && bits.pos + row_budget <= bits.total_bits();
300        if fast {
301            let mut rb = RowBits::start(bits.data, bits.pos);
302            decode_row_fast(&mut rb, t, vp, line, &map);
303            bits.pos = rb.position();
304        } else {
305            decode_row_exact(&mut bits, t, vp, line, &map);
306        }
307        if bits.dead && bits.dead_row.is_none() {
308            bits.dead_row = Some(row);
309        }
310    }
311    Ok(bits)
312}
313
314/// [`decode_predictive`] that entropy-decodes large single-table streams on all threads
315/// (`par_entropy`) when `parallel` is set. Returns the first row that ran past
316/// its data ([`ExactBits::dead_row`]; `None` when the stream was complete). Output is identical to [`decode_predictive`]: streams whose
317/// speculation cannot be proven equal (table switches, truncation, invalid codes) take
318/// the sequential path. `declared` is the payload length the container states (`src`
319/// may run on to the end of the file); speculation covers about that much.
320pub fn decode_predictive_par<M>(
321    src: &[u8],
322    p: &PredictiveParams,
323    out: &mut [u16],
324    stride: usize,
325    map: M,
326    parallel: bool,
327    declared: usize,
328) -> Result<Option<usize>>
329where
330    M: Fn(u16) -> u16 + Sync,
331{
332    let (width, rows) = (p.width, p.rows);
333    if parallel
334        && p.tables.len() == 1
335        && width >= 2
336        && stride >= width
337        && rows > 0
338        && out.len() >= stride * (rows - 1) + width
339        && rows * width >= crate::ljpeg_par::PARALLEL_MIN_SAMPLES
340        && decode_predictive_parallel(src, p, out, stride, &map, declared, None)
341    {
342        return Ok(None);
343    }
344    Ok(decode_predictive(src, p, out, stride, map)?.dead_row())
345}
346
347/// The stream as one difference per symbol.
348struct DiffScan<'a> {
349    src: &'a [u8],
350    table: &'a DirectHuff,
351    span: u64,
352}
353
354impl crate::par_entropy::SymbolStream for DiffScan<'_> {
355    type Pump<'p>
356        = ExactBits<'p>
357    where
358        Self: 'p;
359    const WIDTH: usize = 1;
360    fn cycle(&self) -> usize {
361        1
362    }
363    fn period(&self) -> usize {
364        1
365    }
366    fn total_bits(&self) -> u64 {
367        self.src.len() as u64 * 8
368    }
369    fn span_bits(&self) -> u64 {
370        self.span
371    }
372    fn end_check(&self, pos: u64) -> crate::par_entropy::EndCheck {
373        // Past the end the reader goes dead and yields zeros: leave that to the
374        // sequential decoder.
375        if pos <= self.total_bits() {
376            crate::par_entropy::EndCheck::Valid
377        } else {
378            crate::par_entropy::EndCheck::Unknown
379        }
380    }
381    fn pump_at(&self, bit: u64) -> ExactBits<'_> {
382        ExactBits { data: self.src, pos: bit, dead: false, dead_row: None }
383    }
384    fn position(p: &ExactBits<'_>) -> u64 {
385        p.pos
386    }
387    #[inline(always)]
388    fn decode(&self, p: &mut ExactBits<'_>, _phase: usize, out: &mut [u16]) -> bool {
389        let before = p.pos;
390        out[0] = p.diff(self.table) as u16;
391        if p.pos == before {
392            // A code of length zero (unused table slot): the sequential decoder would not
393            // advance either; treat it as invalid so it decodes this stream itself.
394            p.pos += 1;
395            return false;
396        }
397        true
398    }
399}
400
401fn decode_predictive_parallel<M>(
402    src: &[u8],
403    p: &PredictiveParams,
404    out: &mut [u16],
405    stride: usize,
406    map: &M,
407    declared: usize,
408    chunk: Option<usize>,
409) -> bool
410where
411    M: Fn(u16) -> u16 + Sync,
412{
413    let (width, rows) = (p.width, p.rows);
414    // Declared lengths can be a little short; a true path running past the span falls
415    // back to sequential decoding, so a margin is cheap insurance.
416    let span = if declared == 0 { src.len() } else { (declared + declared / 16 + (64 << 10)).min(src.len()) };
417    let Some(chunk) = chunk.or_else(|| crate::par_entropy::default_chunk(span)) else { return false };
418    let scan = DiffScan { src, table: p.tables[0].1, span: span as u64 * 8 };
419    let Some(values) = crate::par_entropy::decode(&scan, rows * width, chunk) else { return false };
420    // Rows depend on the rows above only through the first two columns of the previous
421    // row of the same parity: chain those, then rebuild rows in parallel.
422    let mut vp = p.vpred;
423    let starts: Vec<[u16; 2]> = (0..rows)
424        .map(|r| {
425            let st = vp[r & 1];
426            for c in 0..2 {
427                vp[r & 1][c] = vp[r & 1][c].wrapping_add(values.get(r * width + c));
428            }
429            st
430        })
431        .collect();
432    let n = out.len().min(rows * stride);
433    flashraw_core::par::for_each_row_chunk(&mut out[..n], stride, 16, true, |r0, chunk| {
434        let mut d = vec![0u16; width];
435        for (i, line) in chunk.chunks_mut(stride).enumerate() {
436            let r = r0 + i;
437            values.copy_to(r * width, &mut d);
438            let mut vpl = starts[r];
439            let mut h = [0u16; 2];
440            for (col, (v, &dv)) in line[..width].iter_mut().zip(&d).enumerate() {
441                let c = col & 1;
442                if col < 2 {
443                    vpl[c] = vpl[c].wrapping_add(dv);
444                    h[c] = vpl[c];
445                } else {
446                    h[c] = h[c].wrapping_add(dv);
447                }
448                *v = map(h[c]);
449            }
450        }
451    });
452    true
453}
454
455/// Register bit cache for rows that are provably inside the data (see the row budget
456/// in [`decode_predictive`]): branchless refills keep 56..=63 bits buffered, so a
457/// symbol and its difference bits never need a second load.
458struct RowBits<'a> {
459    data: &'a [u8],
460    /// Next byte to load.
461    ptr: usize,
462    /// Valid bits, left aligned (zero below them).
463    buf: u64,
464    nbits: u32,
465}
466
467impl<'a> RowBits<'a> {
468    fn start(data: &'a [u8], pos: u64) -> Self {
469        let mut b = RowBits { data, ptr: (pos >> 3) as usize, buf: 0, nbits: 0 };
470        b.refill();
471        let skip = (pos & 7) as u32;
472        b.buf <<= skip;
473        b.nbits -= skip;
474        b
475    }
476
477    /// Bit position in the data.
478    fn position(&self) -> u64 {
479        self.ptr as u64 * 8 - self.nbits as u64
480    }
481
482    #[inline(always)]
483    fn refill(&mut self) {
484        let w = u64::from_be_bytes(self.data[self.ptr..self.ptr + 8].try_into().unwrap());
485        self.buf |= w >> self.nbits;
486        self.ptr += ((63 - self.nbits) >> 3) as usize;
487        self.nbits |= 56;
488    }
489
490    #[inline(always)]
491    fn diff(&mut self, t: &DirectHuff) -> i32 {
492        self.refill();
493        let e = t.fast[(self.buf >> (64 - FAST_BITS)) as usize];
494        if e & FLAG_FULL != 0 {
495            let n = e & 0x1F;
496            self.buf <<= n;
497            self.nbits -= n;
498            return (e >> 16) as u16 as i16 as i32;
499        }
500        self.diff_long(t, e)
501    }
502
503    #[cold]
504    #[inline(never)]
505    fn diff_long(&mut self, t: &DirectHuff, e: u32) -> i32 {
506        let (clen, sym) = if e == SLOW_MARK {
507            let d = t.direct((self.buf >> (64 - t.peek_bits)) as u32);
508            ((d >> 8) as u32, (d & 0xFF) as u32)
509        } else {
510            (e & 0x1F, (e >> 8) & 0xFF)
511        };
512        let n = (sym & 15).saturating_sub(sym >> 4);
513        let w2 = self.buf << clen;
514        let b = if n == 0 { 0 } else { (w2 >> (64 - n)) as u32 };
515        let used = clen + n;
516        self.buf = if used >= 64 { 0 } else { self.buf << used };
517        self.nbits -= used;
518        symbol_diff(b, sym)
519    }
520}
521
522#[inline(always)]
523fn decode_row_fast<M: Fn(u16) -> u16>(
524    bits: &mut RowBits,
525    t: &DirectHuff,
526    vp: &mut [u16; 2],
527    line: &mut [u16],
528    map: &M,
529) {
530    let d0 = bits.diff(t);
531    vp[0] = vp[0].wrapping_add(d0 as u16);
532    let d1 = bits.diff(t);
533    vp[1] = vp[1].wrapping_add(d1 as u16);
534    let (mut h0, mut h1) = (vp[0], vp[1]);
535    line[0] = map(h0);
536    line[1] = map(h1);
537    let mut pairs = line[2..].chunks_exact_mut(2);
538    for pair in &mut pairs {
539        h0 = h0.wrapping_add(bits.diff(t) as u16);
540        h1 = h1.wrapping_add(bits.diff(t) as u16);
541        pair[0] = map(h0);
542        pair[1] = map(h1);
543    }
544    if let [last] = pairs.into_remainder() {
545        h0 = h0.wrapping_add(bits.diff(t) as u16);
546        *last = map(h0);
547    }
548}
549
550#[cold]
551fn decode_row_exact<M: Fn(u16) -> u16>(
552    bits: &mut ExactBits,
553    t: &DirectHuff,
554    vp: &mut [u16; 2],
555    line: &mut [u16],
556    map: &M,
557) {
558    let mut h = [0u16; 2];
559    for (col, v) in line.iter_mut().enumerate() {
560        let d = bits.diff(t) as u16;
561        let c = col & 1;
562        if col < 2 {
563            vp[c] = vp[c].wrapping_add(d);
564            h[c] = vp[c];
565        } else {
566            h[c] = h[c].wrapping_add(d);
567        }
568        *v = map(h[c]);
569    }
570}
571
572/// Nikon Huffman trees (16 counts, then symbols).
573pub const NIKON_TREES: [&[u8]; 6] = [
574    // 12-bit lossy
575    // (the 14th code maps to symbol 0: the table is zero padded)
576    &[0, 1, 5, 1, 1, 1, 1, 1, 1, 2, 0, 0, 0, 0, 0, 0, 5, 4, 3, 6, 2, 7, 1, 0, 8, 9, 11, 10, 12, 0],
577    // 12-bit lossy after split
578    &[0, 1, 5, 1, 1, 1, 1, 1, 1, 2, 0, 0, 0, 0, 0, 0, 0x39, 0x5a, 0x38, 0x27, 0x16, 5, 4, 3, 2, 1, 0, 11, 12, 12],
579    // 12-bit lossless
580    &[0, 1, 4, 2, 3, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 4, 6, 3, 7, 2, 8, 1, 9, 0, 10, 11, 12],
581    // 14-bit lossy
582    &[0, 1, 4, 3, 1, 1, 1, 1, 1, 2, 0, 0, 0, 0, 0, 0, 5, 6, 4, 7, 8, 3, 9, 2, 1, 0, 10, 11, 12, 13, 14],
583    // 14-bit lossy after split
584    &[0, 1, 5, 1, 1, 1, 1, 1, 1, 1, 2, 0, 0, 0, 0, 0, 8, 0x5c, 0x4b, 0x3a, 0x29, 7, 6, 5, 4, 3, 2, 1, 0, 13, 14],
585    // 14-bit lossless
586    &[0, 1, 4, 2, 2, 3, 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 7, 6, 8, 5, 9, 4, 10, 3, 11, 12, 2, 0, 1, 13, 14],
587];
588
589/// Decoded Nikon compression metadata (makernote tag 0x96 "NEFLinearizationTable").
590#[derive(Debug, Clone)]
591pub struct NikonMeta {
592    /// Version bytes (0x44 0x10/0x20/0x40 lossy, 0x46 0x30 lossless, ...).
593    pub version: [u8; 2],
594    pub vpred: [[u16; 2]; 2],
595    /// Index into [`NIKON_TREES`].
596    pub tree: usize,
597    /// Row at which the "after split" tree takes over (0 = none).
598    pub split_row: usize,
599    /// Linearisation curve indexed by the clamped predictor value (0x4000 entries).
600    pub curve: Vec<u16>,
601}
602
603impl NikonMeta {
604    /// Parse the blob at makernote tag 0x96 (or 0x8C). `meta` starts at the tag's value
605    /// and may extend to the end of the file (parsing may read beyond the tag).
606    pub fn parse(meta: &[u8], order: ByteOrder, bits: u32) -> Result<Self> {
607        if bits != 12 && bits != 14 {
608            return Err(Error::unsupported_compression("NEF", format!("{bits}-bit Nikon Huffman")));
609        }
610        let rd = |off: usize| -> Result<u16> {
611            let b = meta.get(off..off + 2).ok_or_else(|| Error::corrupt("truncated Nikon linearization table"))?;
612            Ok(order.u16([b[0], b[1]]))
613        };
614        let ver0 = *meta.first().ok_or_else(|| Error::corrupt("empty Nikon linearization table"))?;
615        let ver1 = *meta.get(1).ok_or_else(|| Error::corrupt("truncated Nikon linearization table"))?;
616        let mut pos = 2usize;
617        if ver0 == 0x49 || ver1 == 0x58 {
618            pos += 2110;
619        }
620        let mut vpred = [[0u16; 2]; 2];
621        for (i, v) in vpred.iter_mut().flatten().enumerate() {
622            *v = rd(pos + 2 * i)?;
623        }
624        pos += 8;
625        let mut tree = if ver0 == 0x46 { 2 } else { 0 };
626        if bits == 14 {
627            tree += 3;
628        }
629        // Linearisation curve (identity unless the table provides one).
630        let mut curve: Vec<u16> = (0..=0xFFFFu32).map(|i| i as u16).collect();
631        let mut max = (1usize << bits) & 0x7fff;
632        let mut step = max;
633        let csize = rd(pos)? as usize;
634        pos += 2;
635        if csize > 1 {
636            step = max / (csize - 1);
637        }
638        if ver0 == 0x44 && (ver1 == 0x20 || (ver1 == 0x40 && step > 3)) && step > 0 {
639            if ver1 == 0x40 {
640                step /= 4;
641                max /= 4;
642            }
643            for i in 0..csize {
644                let idx = i * step;
645                let v = rd(pos + 2 * i)?;
646                if let Some(c) = curve.get_mut(idx) {
647                    *c = v;
648                }
649            }
650            if step > 0 {
651                for i in 0..max {
652                    let base = i - i % step;
653                    let (a, b) = (curve[base] as u64, curve.get(base + step).copied().unwrap_or(0) as u64);
654                    curve[i] = ((a * (step - i % step) as u64 + b * (i % step) as u64) / step as u64) as u16;
655                }
656            }
657        } else if ver0 != 0x46 && csize <= 0x4001 {
658            for (i, c) in curve.iter_mut().enumerate().take(csize) {
659                *c = rd(pos + 2 * i)?;
660            }
661        }
662        let split_row = if ver0 == 0x44 && (ver1 == 0x20 || ver1 == 0x40) { rd(562)? as usize } else { 0 };
663        curve.truncate(0x4000);
664        Ok(NikonMeta { version: [ver0, ver1], vpred, tree, split_row, curve })
665    }
666
667    /// Whether this is a lossless stream (no curve, identity mapping).
668    pub fn is_lossless(&self) -> bool {
669        self.version[0] == 0x46
670    }
671
672    /// Largest value the curve produces (the effective white level of lossy files).
673    pub fn curve_max(&self) -> u16 {
674        self.curve.iter().copied().max().unwrap_or(0)
675    }
676}
677
678/// Decode a Nikon Huffman NEF payload. `src` should extend to the end of the file
679/// (decoding reads past the declared strip length).
680#[allow(clippy::too_many_arguments)]
681pub fn decode_nikon(
682    src: &[u8],
683    meta: &NikonMeta,
684    width: usize,
685    rows: usize,
686    out: &mut [u16],
687    stride: usize,
688    parallel: bool,
689    declared: usize,
690) -> Result<()> {
691    let dead_row = decode_nikon_rows(src, meta, width, rows, out, stride, parallel, declared)?;
692    check_complete(dead_row, rows)
693}
694
695/// Like [`decode_nikon`], but a stream that runs out of data is not an error: returns the
696/// row in which the data ended (`None` when the stream was complete). That row and the
697/// rows after it hold no real data (callers recovering a truncated file zero them).
698#[allow(clippy::too_many_arguments)]
699pub fn decode_nikon_rows(
700    src: &[u8],
701    meta: &NikonMeta,
702    width: usize,
703    rows: usize,
704    out: &mut [u16],
705    stride: usize,
706    parallel: bool,
707    declared: usize,
708) -> Result<Option<usize>> {
709    let main = DirectHuff::from_counts(NIKON_TREES[meta.tree])?;
710    let after;
711    let mut tables: Vec<(usize, &DirectHuff)> = vec![(0, &main)];
712    if meta.split_row > 0 && meta.tree + 1 < NIKON_TREES.len() {
713        after = DirectHuff::from_counts(NIKON_TREES[meta.tree + 1])?;
714        tables.push((meta.split_row, &after));
715    }
716    let p = PredictiveParams { width, rows, tables: &tables, vpred: meta.vpred };
717    let curve: &[u16; 0x4000] =
718        meta.curve.as_slice().try_into().map_err(|_| Error::corrupt("Nikon curve has the wrong size"))?;
719    let dead_row = decode_predictive_par(
720        src,
721        &p,
722        out,
723        stride,
724        |h| {
725            // Clamp the signed predictor into the curve's index range.
726            let i = (h as i16).clamp(0, 0x3fff) as usize;
727            curve[i]
728        },
729        parallel,
730        declared,
731    )?;
732    Ok(dead_row)
733}
734
735/// A stream may need a few zero padding bits for its very last codes (encoders do not
736/// always flush), but running dry earlier means the file is truncated: the remaining rows
737/// would be flat garbage, so report it instead of returning a "successful" image.
738pub fn check_complete(dead_row: Option<usize>, rows: usize) -> Result<()> {
739    match dead_row {
740        Some(r) if r + 1 < rows => {
741            Err(Error::corrupt(format!("compressed data ends at row {r}/{rows} (truncated file)")))
742        }
743        _ => Ok(()),
744    }
745}
746
747#[cfg(test)]
748mod tests {
749    use super::*;
750
751    /// Bit writer for synthetic streams.
752    struct W {
753        out: Vec<u8>,
754        acc: u64,
755        n: u32,
756    }
757    impl W {
758        fn new() -> Self {
759            W { out: vec![], acc: 0, n: 0 }
760        }
761        fn put(&mut self, v: u32, len: u32) {
762            for i in (0..len).rev() {
763                self.acc = (self.acc << 1) | ((v >> i) & 1) as u64;
764                self.n += 1;
765                if self.n == 8 {
766                    self.out.push(self.acc as u8);
767                    self.acc = 0;
768                    self.n = 0;
769                }
770            }
771        }
772        fn finish(mut self) -> Vec<u8> {
773            while self.n != 0 {
774                self.put(0, 1);
775            }
776            self.out
777        }
778    }
779
780    /// Canonical codes for a counts+symbols source table: (symbol, code, length).
781    fn codes(src: &[u8]) -> Vec<(u8, u32, u32)> {
782        let mut out = vec![];
783        let mut code = 0u32;
784        let mut k = 16;
785        for len in 1..=16u32 {
786            for _ in 0..src[len as usize - 1] {
787                out.push((src[k], code, len));
788                code += 1;
789                k += 1;
790            }
791            code <<= 1;
792        }
793        out
794    }
795
796    /// Encode a diff for a shl==0 symbol table.
797    fn encode(w: &mut W, cs: &[(u8, u32, u32)], d: i32) {
798        let len = if d == 0 { 0 } else { 32 - d.unsigned_abs().leading_zeros() };
799        let (_, c, l) = cs.iter().find(|(s, _, _)| *s as u32 == len).expect("symbol");
800        w.put(*c, *l);
801        if len > 0 {
802            let v = if d < 0 { (d - 1) as u32 & ((1 << len) - 1) } else { d as u32 };
803            w.put(v, len);
804        }
805    }
806
807    fn reference_decode(src: &[u8], t: &DirectHuff, width: usize, rows: usize, vpred: [[u16; 2]; 2]) -> Vec<u16> {
808        let mut out = vec![0u16; width * rows];
809        let mut vp = vpred;
810        let mut b = ExactBits::new(src);
811        for row in 0..rows {
812            decode_row_exact(&mut b, t, &mut vp[row & 1], &mut out[row * width..(row + 1) * width], &|h| h);
813        }
814        out
815    }
816
817    #[test]
818    fn lossless_roundtrip_all_trees() {
819        for &tree in &[2usize, 5] {
820            let src = NIKON_TREES[tree];
821            let t = DirectHuff::from_counts(src).unwrap();
822            let cs = codes(src);
823            let (width, rows) = (37usize, 9usize);
824            let maxv = if tree == 2 { 4095 } else { 16383 };
825            let img: Vec<u16> = (0..width * rows).map(|i| ((i * 7919 + i * i * 31) % (maxv + 1)) as u16).collect();
826            let vpred = [[100u16, 200], [300, 400]];
827            let mut w = W::new();
828            let mut vp = vpred;
829            for row in 0..rows {
830                let mut h = [0u16; 2];
831                for col in 0..width {
832                    let v = img[row * width + col];
833                    let c = col & 1;
834                    let pred = if col < 2 { vp[row & 1][c] } else { h[c] };
835                    let d = v as i32 - pred as i32;
836                    encode(&mut w, &cs, d);
837                    if col < 2 {
838                        vp[row & 1][c] = v;
839                    }
840                    h[c] = v;
841                }
842            }
843            let mut data = w.finish();
844            data.extend_from_slice(&[0u8; 4096]);
845            let mut out = vec![0u16; width * rows];
846            let p = PredictiveParams { width, rows, tables: &[(0, &t)], vpred };
847            decode_predictive(&data, &p, &mut out, width, |h| h).unwrap();
848            assert_eq!(out, img, "tree {tree}");
849            assert_eq!(reference_decode(&data, &t, width, rows, vpred), img);
850        }
851    }
852
853    #[test]
854    fn fast_path_matches_exact_path_on_random_streams() {
855        let mut x = 0x1234_5678_9abc_def0u64;
856        let data: Vec<u8> = (0..200_000)
857            .map(|_| {
858                x ^= x << 13;
859                x ^= x >> 7;
860                x ^= x << 17;
861                x as u8
862            })
863            .collect();
864        for src in NIKON_TREES {
865            let t = DirectHuff::from_counts(src).unwrap();
866            let (width, rows) = (301, 40);
867            let vpred = [[1, 2], [3, 4]];
868            let mut fast = vec![0u16; width * rows];
869            let p = PredictiveParams { width, rows, tables: &[(0, &t)], vpred };
870            decode_predictive(&data, &p, &mut fast, width, |h| h).unwrap();
871            assert_eq!(fast, reference_decode(&data, &t, width, rows, vpred));
872        }
873    }
874
875    #[test]
876    fn truncation_freezes_predictors() {
877        let t = DirectHuff::from_counts(NIKON_TREES[2]).unwrap();
878        let data = [0xA5u8; 10];
879        let (width, rows) = (16, 8);
880        let mut out = vec![0u16; width * rows];
881        let p = PredictiveParams { width, rows, tables: &[(0, &t)], vpred: [[0, 0], [0, 0]] };
882        let b = decode_predictive(&data, &p, &mut out, width, |h| h).unwrap();
883        assert!(b.is_dead());
884        // After death, rows repeat the frozen vertical predictors and differences are zero.
885        let last = &out[(rows - 1) * width..];
886        assert!(last.chunks(2).all(|c| c == &last[..2]));
887    }
888
889    #[test]
890    fn truncated_streams_are_corrupt() {
891        let t = DirectHuff::from_counts(NIKON_TREES[2]).unwrap();
892        let data = [0xA5u8; 10];
893        let (width, rows) = (16, 8);
894        let mut out = vec![0u16; width * rows];
895        let p = PredictiveParams { width, rows, tables: &[(0, &t)], vpred: [[0, 0], [0, 0]] };
896        let b = decode_predictive(&data, &p, &mut out, width, |h| h).unwrap();
897        let dead = b.dead_row().expect("ran past the data");
898        assert!(dead + 1 < rows);
899        assert!(matches!(check_complete(Some(dead), rows), Err(Error::CorruptFile { .. })));
900        // Needing padding bits only for the last row is tolerated (unflushed encoders).
901        assert!(check_complete(Some(rows - 1), rows).is_ok());
902        assert!(check_complete(None, rows).is_ok());
903    }
904
905    #[test]
906    fn shl_symbols() {
907        // len 9, shl 3: 6 transmitted bits.
908        let sym = 0x39;
909        assert_eq!(symbol_diff(0b111111, sym), (0b111111 << 3) | 0b100);
910        let d = symbol_diff(0, sym);
911        assert!(d < 0);
912    }
913
914    #[test]
915    fn meta_parsing_rejects_garbage() {
916        assert!(NikonMeta::parse(&[], ByteOrder::Big, 12).is_err());
917        assert!(NikonMeta::parse(&[0x44, 0x20, 0, 0], ByteOrder::Big, 12).is_err());
918        assert!(NikonMeta::parse(&[0x46, 0x30, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], ByteOrder::Big, 10).is_err());
919        let m = NikonMeta::parse(&[0x46, 0x30, 0, 1, 0, 2, 0, 3, 0, 4, 0, 0], ByteOrder::Big, 14).unwrap();
920        assert_eq!(m.vpred, [[1, 2], [3, 4]]);
921        assert_eq!(m.tree, 5);
922        assert!(m.is_lossless());
923        assert_eq!(m.curve[1234], 1234);
924    }
925
926    #[test]
927    fn parallel_matches_sequential() {
928        let t = DirectHuff::from_counts(NIKON_TREES[2]).unwrap();
929        let (width, rows) = (96usize, 50usize);
930        let mut x = 99u32;
931        let data: Vec<u8> = (0..width * rows * 2)
932            .map(|_| {
933                x = x.wrapping_mul(1103515245).wrapping_add(12345);
934                (x >> 24) as u8
935            })
936            .collect();
937        let p = PredictiveParams { width, rows, tables: &[(0, &t)], vpred: [[100, 200], [300, 400]] };
938        let stride = width + 3;
939        let mut seq = vec![0u16; stride * rows];
940        decode_predictive(&data, &p, &mut seq, stride, |h| h ^ 0x55).unwrap();
941        for chunk in [3usize, 50, 999] {
942            let mut par = vec![0u16; stride * rows];
943            assert!(decode_predictive_parallel(&data, &p, &mut par, stride, &|h| h ^ 0x55, 0, Some(chunk)), "{chunk}");
944            assert_eq!(par, seq, "chunk {chunk}");
945        }
946    }
947}