wai/wai-rs/src/jpegai_bitstream.rs

@@ -11,6 +11,12 @@
 //! bit-granular cursor. Reads past the end return the bits available (zero-extended),
 //! matching the reference's bounds check.
 
+/// Bits needed to carry `max_value`, per the reference `HeaderCoder::calc_bits_count`:
+/// `ceil(log2(max_value + 1))`. `max_value = 0` needs no bits.
+pub fn bits_for_max(max_value: u64) -> u32 {
+    if max_value == 0 { 0 } else { u64::BITS - max_value.leading_zeros() }
+}
+
 /// Marker identifiers (big-endian 2-byte prefixes).
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum Marker {
@@ -90,6 +96,16 @@ impl<'a> BitReader<'a> {
         v
     }
 
+    /// Read one symbol whose width is **derived from its maximum value**, per the
+    /// reference `HeaderCoder::calc_bits_count`: `bits = ceil(log2(max_value + 1))`.
+    ///
+    /// This is the header's universal width rule — no field has a fixed width. Some are
+    /// runtime-dependent: a `max_symbol` decoded earlier in a segment sets the width of
+    /// the weights that follow it (e.g. 65535 ⇒ 16 bits, then 99 ⇒ 7 bits).
+    pub fn read_max(&mut self, max_value: u64) -> u64 {
+        self.read(bits_for_max(max_value))
+    }
+
     /// Unsigned exp-Golomb, k=0: count leading zeros `n`, then read `n` bits;
     /// value = `2^n − 1 + suffix`. Used for marker-segment payload lengths.
     pub fn read_expgolomb_k0(&mut self) -> u64 {
@@ -196,6 +212,205 @@ pub struct PictureHeader {
 /// Branches not exercised by the 4:4:4 single-tile reference stream (region
 /// sub-splits, `independent_beta_uv`'s second β, the signalled colour matrix,
 /// multi-threading counts) follow the reference syntax and are decoded but untested.
+// ---- tools header (TON) ----
+
+/// One `decode_filters` group of the EFE **linear** filter header.
+#[derive(Clone, Debug, Default)]
+pub struct EflFilters {
+    /// `best_cand_idx − 1`; `-1` means the plane's filter is absent.
+    pub cand_u: i32,
+    pub cand_v: i32,
+    pub fl_u: u32,
+    pub fl_v: u32,
+    pub min_symbol: u32,
+    pub max_symbol: u32,
+    pub weights_u: Vec<u32>,
+    pub weights_u2: Vec<u32>,
+    pub weights_v: Vec<u32>,
+    pub weights_v2: Vec<u32>,
+}
+
+/// EFE linear filter header: the two plane means plus two filter groups.
+#[derive(Clone, Debug, Default)]
+pub struct EflHeader {
+    /// `B1[i] / 100` = the per-plane mean consumed by `jpegai_pipeline::efl_chroma`.
+    pub b1: [u32; 2],
+    pub filters_best2: EflFilters,
+    pub filters_best: EflFilters,
+}
+
+/// eICCI per-tile model selection.
+#[derive(Clone, Debug, Default)]
+pub struct IcciTile {
+    /// Per-plane use flags (Y, U, V).
+    pub use_yuv: [bool; 3],
+    pub use_short_list: bool,
+    /// Signalled model indices (into the active list) for the Y and the joint-UV paths.
+    pub model_idx_y: u32,
+    pub model_idx_uv: u32,
+}
+
+/// EFE nonlinear filter header (the luma-aided chroma filter's parameters).
+#[derive(Clone, Debug, Default)]
+pub struct EfnHeader {
+    pub min_symbol: u32,
+    pub max_symbol: u32,
+    pub mask1_enabled: bool,
+    pub mask2_enabled: bool,
+    pub nonlinear_u_enabled: bool,
+    pub nonlinear_v_enabled: bool,
+    pub tile_width: u32,
+    pub tile_height: u32,
+    /// Per-tile luma range — the `luma_min_q` / `luma_max_q` of `efn_filter`.
+    pub luma_min: Vec<u32>,
+    pub luma_max: Vec<u32>,
+    /// Per-plane filter weights (`A1`), already offset by `min_symbol` — the
+    /// `weights_q15` of `efn_filter`.
+    pub weights_u: Vec<u32>,
+    pub weights_v: Vec<u32>,
+}
+
+/// The decoded tools header (TON substream).
+#[derive(Clone, Debug, Default)]
+pub struct ToolsHeader {
+    pub efl_enabled: bool,
+    pub efl: Option<EflHeader>,
+    pub icci_enabled: bool,
+    pub icci_tile_manager_enabled: bool,
+    pub icci_tiles: Vec<IcciTile>,
+    pub efn_enabled: bool,
+    pub efn: Option<EfnHeader>,
+    pub lef_enabled: bool,
+    /// LEF's reference channel index.
+    pub lef_ch_idx: u32,
+    /// Bit position just past the last field (before byte-alignment padding).
+    pub end_bit: usize,
+}
+
+/// `cands[i][0]` — the spatial-tiling candidate filter counts of the EFE linear filter.
+const EFL_CAND_FILTER_COUNT: [u32; 6] = [1, 2, 2, 3, 3, 4];
+
+/// One `decode_filters` group. `scale_2x2` selects the 4:4:4 path, where a single
+/// plane's coefficients are decoded and replicated rather than coded four times.
+fn parse_efl_filters(br: &mut BitReader, scale_2x2: bool) -> EflFilters {
+    let mut f = EflFilters::default();
+    f.cand_u = br.read_max(15) as i32 - 1;
+    f.cand_v = br.read_max(15) as i32 - 1;
+    if f.cand_u > -1 { f.fl_u = br.read_max(9) as u32; }
+    if f.cand_v > -1 { f.fl_v = br.read_max(9) as u32; }
+    f.min_symbol = br.read_max((1 << 16) - 1) as u32;
+    f.max_symbol = br.read_max((1 << 16) - 1) as u32;
+    let mut group = |cand: i32, fl: u32, br: &mut BitReader| -> (Vec<u32>, Vec<u32>) {
+        if cand <= -1 { return (Vec::new(), Vec::new()); }
+        let taps = (fl * fl) as usize;
+        // shape[0]==4 is coded as a single plane under 4:4:4 and replicated
+        let n_main = if scale_2x2 { taps } else { EFL_CAND_FILTER_COUNT[cand as usize] as usize * taps };
+        let main = (0..n_main).map(|_| br.read_max(f.max_symbol as u64) as u32).collect();
+        let second = (0..taps).map(|_| br.read_max(f.max_symbol as u64) as u32).collect();
+        (main, second)
+    };
+    let (u, u2) = group(f.cand_u, f.fl_u, br);
+    f.weights_u = u;
+    f.weights_u2 = u2;
+    let (v, v2) = group(f.cand_v, f.fl_v, br);
+    f.weights_v = v;
+    f.weights_v2 = v2;
+    f
+}
+
+/// Parse the **post-filter chain** of the tools header (TON), in filters-factory order:
+/// EFE-linear → ICCI → eICCI → EFE-nonlinear → LEF. Each tool is gated by its own enable
+/// flag (the recursive header walk), and every field's width is derived from its
+/// `max_symbol_value` — including the runtime-dependent ones, where a just-decoded
+/// `max_symbol` sets the width of the weights that follow.
+///
+/// `start_bit` is the bit position of the **first filter's enable flag**, not the TON
+/// payload start: `_stream_header_part` defaults to `"tool_header"`, so tools other than
+/// the post-filters also emit into TON ahead of the filter chain (in the reference
+/// stream two such bits precede EFE-linear's flag). Locating that prefix generally means
+/// walking the whole tool tree; this function parses the chain from a given offset.
+///
+/// `models_count` is the size of the eICCI model list the stream selects from (it sets
+/// the width of the signalled model indices); `num_tiles` is the eICCI tile count.
+pub fn parse_tools_header(data: &[u8], start_bit: usize, scale_2x2: bool, models_count: u32, num_tiles: usize) -> ToolsHeader {
+    let mut br = BitReader::at(data, start_bit);
+    let mut t = ToolsHeader::default();
+
+    t.efl_enabled = br.read_max(1) == 1;
+    if t.efl_enabled {
+        let mut h = EflHeader::default();
+        h.b1[0] = br.read_max((1 << 15) - 1) as u32;
+        h.b1[1] = br.read_max((1 << 15) - 1) as u32;
+        h.filters_best2 = parse_efl_filters(&mut br, scale_2x2);
+        h.filters_best = parse_efl_filters(&mut br, scale_2x2);
+        t.efl = Some(h);
+    }
+
+    t.icci_enabled = br.read_max(1) == 1;
+    if t.icci_enabled {
+        t.icci_tile_manager_enabled = br.read_max(1) == 1;
+        for _ in 0..num_tiles.max(1) {
+            let mut tile = IcciTile::default();
+            for u in tile.use_yuv.iter_mut() { *u = br.read_max(1) == 1; }
+            if tile.use_yuv.iter().any(|&b| b) {
+                tile.use_short_list = br.read_max(1) == 1;
+                if tile.use_yuv[0] {
+                    tile.model_idx_y = br.read_max(models_count as u64) as u32;
+                }
+                if tile.use_yuv[1] || tile.use_yuv[2] {
+                    tile.model_idx_uv = br.read_max(models_count as u64) as u32;
+                }
+            }
+            t.icci_tiles.push(tile);
+        }
+    }
+
+    t.efn_enabled = br.read_max(1) == 1;
+    if t.efn_enabled {
+        let mut e = EfnHeader::default();
+        e.min_symbol = br.read_max((1 << 16) - 1) as u32;
+        e.max_symbol = br.read_max((1 << 16) - 1) as u32;
+        e.mask1_enabled = br.read_max(1) == 1;
+        e.mask2_enabled = br.read_max(1) == 1;
+        if e.mask1_enabled || e.mask2_enabled {
+            let _b_size = br.read_max((1 << 10) - 1);
+            let h = br.read_max((1 << 10) - 1) as usize;
+            let w = br.read_max((1 << 10) - 1) as usize;
+            // W5 masks: only the enabled one(s) carry h×w symbols (max 2 ⇒ 2 bits)
+            for enabled in [e.mask1_enabled, e.mask2_enabled] {
+                let n = if enabled { h * w } else { 0 };
+                for _ in 0..n { br.read_max(2); }
+            }
+        }
+        e.nonlinear_u_enabled = br.read_max(1) == 1;
+        e.nonlinear_v_enabled = br.read_max(1) == 1;
+        if e.nonlinear_u_enabled || e.nonlinear_v_enabled {
+            e.tile_width = br.read_max((1 << 16) - 1) as u32;
+            e.tile_height = br.read_max((1 << 16) - 1) as u32;
+            let n = br.read_max((1 << 16) - 1) as usize; // numTiles
+            e.luma_min = (0..n).map(|_| br.read_max((1 << 16) - 1) as u32).collect();
+            e.luma_max = (0..n).map(|_| br.read_max((1 << 16) - 1) as u32).collect();
+            // candNum + A1 are coded per ENABLED plane (U first, then V)
+            for (enabled, dst) in [(e.nonlinear_u_enabled, 0), (e.nonlinear_v_enabled, 1)] {
+                if !enabled { continue; }
+                let cand_num = br.read_max((1 << 16) - 1) as usize;
+                let w: Vec<u32> = (0..cand_num)
+                    .map(|_| br.read_max(e.max_symbol as u64) as u32 + e.min_symbol)
+                    .collect();
+                if dst == 0 { e.weights_u = w; } else { e.weights_v = w; }
+            }
+        }
+        t.efn = Some(e);
+    }
+
+    t.lef_enabled = br.read_max(1) == 1;
+    if t.lef_enabled {
+        t.lef_ch_idx = br.read_max(255) as u32;
+    }
+    t.end_bit = br.bit_pos();
+    t
+}
+
 pub fn parse_picture_header(data: &[u8], start_bit: usize) -> PictureHeader {
     let mut br = BitReader::at(data, start_bit);
     let mut h = PictureHeader::default();

wai/wai-rs/tests/jpegai_bitstream_calibration.rs

@@ -9,7 +9,7 @@
 //! env: WAI_JPEGAI_CODESTREAM = path to test256.bin
 #![cfg(feature = "neural_int")]
 
-use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, BitReader, Marker};
+use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, parse_tools_header, BitReader, Marker};
 
 fn codestream() -> Option<Vec<u8>> {
     match std::env::var("WAI_JPEGAI_CODESTREAM") {
@@ -126,3 +126,76 @@ fn full_picture_header_matches_reference() {
     assert_eq!(h.end_bit, 437, "all PIH fields consumed");
     assert!(h.end_bit <= (pih.start + pih.len) * 8, "within the PIH payload");
 }
+
+#[test]
+fn tools_header_matches_reference() {
+    let Some(data) = codestream() else { return };
+    let segs = parse_container(&data);
+    let ton = find_segment(&segs, Marker::Ton).expect("TON present");
+    let pih = parse_picture_header(&data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
+
+    // 4:4:4 ⇒ scale_ver = scale_hor = 3 − c/s = 2 ⇒ the single-plane filter path
+    let scale_2x2 = (3 - pih.c_ver / pih.s_ver == 2) && (3 - pih.c_hor / pih.s_hor == 2);
+    assert!(scale_2x2, "test256 is 4:4:4");
+
+    // The post-filter chain starts 2 bits into TON: `_stream_header_part` defaults to
+    // "tool_header", so other tools emit into TON ahead of the filters. models_count
+    // sets the signalled-index width (2 bits here).
+    let filters_bit = ton.start_bit() + 2;
+    let t = parse_tools_header(&data, filters_bit, scale_2x2, 3, 1);
+    eprintln!("{t:#?}");
+
+    // EFE linear: the two plane means feeding efl_chroma
+    assert!(t.efl_enabled);
+    let efl = t.efl.as_ref().unwrap();
+    assert_eq!(efl.b1, [12561, 13091], "B1/100 = 125.61 / 130.91 = the EFL plane means");
+    let f1 = &efl.filters_best2;
+    assert_eq!((f1.cand_u, f1.cand_v), (0, 0));
+    assert_eq!((f1.fl_u, f1.fl_v), (1, 1));
+    assert_eq!((f1.min_symbol, f1.max_symbol), (0, 65535));
+    assert_eq!((&f1.weights_u[..], &f1.weights_u2[..]), (&[32677u32][..], &[32772u32][..]));
+    assert_eq!((&f1.weights_v[..], &f1.weights_v2[..]), (&[32758u32][..], &[32776u32][..]));
+    let f2 = &efl.filters_best;
+    // the second group's max_symbol (99) narrows the weights to 7 bits
+    assert_eq!((f2.min_symbol, f2.max_symbol), (32677, 99));
+    assert_eq!((&f2.weights_u[..], &f2.weights_u2[..]), (&[0u32][..], &[95u32][..]));
+    assert_eq!((&f2.weights_v[..], &f2.weights_v2[..]), (&[81u32][..], &[99u32][..]));
+
+    // eICCI: the per-tile model selection (the last input the eICCI net needed)
+    assert!(t.icci_enabled);
+    assert!(!t.icci_tile_manager_enabled);
+    assert_eq!(t.icci_tiles.len(), 1);
+    let tile = &t.icci_tiles[0];
+    assert_eq!(tile.use_yuv, [true, true, true]);
+    assert!(tile.use_short_list);
+    assert_eq!((tile.model_idx_y, tile.model_idx_uv), (0, 1));
+
+    // EFE nonlinear: exactly the inputs efn_filter consumes
+    assert!(t.efn_enabled);
+    let efn = t.efn.as_ref().unwrap();
+    assert_eq!((efn.min_symbol, efn.max_symbol), (0, 65535));
+    assert!(!efn.mask1_enabled && !efn.mask2_enabled);
+    assert!(!efn.nonlinear_u_enabled && efn.nonlinear_v_enabled, "only V is filtered");
+    assert_eq!((efn.tile_width, efn.tile_height), (256, 256));
+    assert_eq!((&efn.luma_min[..], &efn.luma_max[..]), (&[497u32][..], &[25320u32][..]),
+               "== efn_fixture luma_min_q / luma_max_q");
+    assert!(efn.weights_u.is_empty(), "U disabled ⇒ no candNum/A1 for U");
+    assert_eq!(efn.weights_v, vec![32966, 32443, 33031, 32361, 33068, 33006, 31968, 33003],
+               "== efn_fixture weights_q15");
+
+    assert!(t.lef_enabled);
+    assert_eq!(t.lef_ch_idx, 0);
+
+    // the whole tools header consumed exactly, up to the pad closing the 63-byte TON
+    assert_eq!(t.end_bit, 973, "all TON fields consumed");
+    assert!(t.end_bit <= (ton.start + ton.len) * 8);
+}
+
+#[test]
+fn header_widths_follow_calc_bits_count() {
+    use wai::jpegai_bitstream::bits_for_max;
+    // bits = ceil(log2(max+1)) — the reference HeaderCoder::calc_bits_count
+    for (max, want) in [(0u64, 0u32), (1, 1), (2, 2), (3, 2), (4, 3), (9, 4), (15, 4), (99, 7), (255, 8), (65535, 16)] {
+        assert_eq!(bits_for_max(max), want, "bits_for_max({max})");
+    }
+}