open-standards/wai/wai-rs/src/jpegai_bitstream.rs

//! JPEG-AI **bitstream container + header bit reader** — the decoder's entry point.
//!
//! A codestream is a sequence of marker segments. Each (except the marker-only SOC/EOC)
//! is `marker(2 B, big-endian) + expgolomb_k0(payload_size) + byte-align + payload`.
//! The header substreams (PIH, TON) are **plain bit-packed**, read MSB-first — the
//! headers are *not* arithmetic-coded; only the latent substreams (SORP/SORS/SOZ) are,
//! and those are handed to [`crate::jpegai_mans`].
//!
//! [`BitReader`] is a faithful port of the reference `EcLibDirect::read_bits`: bits are
//! taken MSB-first within each byte and accumulated MSB-first into the value, with a
//! bit-granular cursor. Reads past the end return the bits available (zero-extended),
//! matching the reference's bounds check.

// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
#![deny(clippy::disallowed_methods)]

/// 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 {
    Soc,  // 0xFF80 start of codestream (no payload)
    Eoc,  // 0xFF81 end of codestream (no payload)
    Pih,  // 0xFF82 picture header
    Ton,  // 0xFF83 tools header
    Rdi,  // 0xFF84
    Soz,  // 0xFF88 hyper-latent substream
    Sorp, // 0xFF89 primary residual substream
    Sors, // 0xFF8A secondary residual substream
    Soq,  // 0xFF8B
    Udi,  // 0xFF8C
}

impl Marker {
    pub fn from_u16(v: u16) -> Option<Marker> {
        Some(match v {
            0xFF80 => Marker::Soc,
            0xFF81 => Marker::Eoc,
            0xFF82 => Marker::Pih,
            0xFF83 => Marker::Ton,
            0xFF84 => Marker::Rdi,
            0xFF88 => Marker::Soz,
            0xFF89 => Marker::Sorp,
            0xFF8A => Marker::Sors,
            0xFF8B => Marker::Soq,
            0xFF8C => Marker::Udi,
            _ => return None,
        })
    }
    /// SOC/EOC carry no length or payload.
    pub fn has_payload(self) -> bool {
        !matches!(self, Marker::Soc | Marker::Eoc)
    }
}

/// One parsed marker segment: where its payload lives in the codestream.
#[derive(Clone, Copy, Debug)]
pub struct Segment {
    pub marker: Marker,
    pub start: usize, // payload start (byte offset)
    pub len: usize,   // payload length in bytes
}

impl Segment {
    /// Bit offset of the payload start — the cursor for a [`BitReader`].
    pub fn start_bit(&self) -> usize { self.start * 8 }
}

/// MSB-first bit reader over a byte buffer, with a bit-granular cursor.
/// Port of the reference `EcLibDirect::read_bits`.
pub struct BitReader<'a> {
    data: &'a [u8],
    pos: usize, // in bits
}

impl<'a> BitReader<'a> {
    pub fn new(data: &'a [u8]) -> Self { BitReader { data, pos: 0 } }
    pub fn at(data: &'a [u8], bit_pos: usize) -> Self { BitReader { data, pos: bit_pos } }
    pub fn bit_pos(&self) -> usize { self.pos }
    pub fn seek_bit(&mut self, bit_pos: usize) { self.pos = bit_pos; }
    /// Advance to the next byte boundary (the reference flushes after the segment length).
    pub fn align_to_byte(&mut self) { self.pos = (self.pos + 7) & !7; }

    /// Read `n` bits MSB-first. Stops early at end-of-buffer (returning what was read),
    /// mirroring the reference's `byte_pos < m_data_size` guard.
    pub fn read(&mut self, n: u32) -> u64 {
        let mut v = 0u64;
        for _ in 0..n {
            let byte = self.pos >> 3;
            if byte >= self.data.len() { break; }
            let bit = (self.data[byte] >> (7 - (self.pos & 7))) & 1;
            v = (v << 1) | bit as u64;
            self.pos += 1;
        }
        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 {
        let mut n = 0u32;
        while self.read(1) == 0 {
            n += 1;
            if n >= 32 || (self.pos >> 3) >= self.data.len() { return 0; }
        }
        let suffix = if n > 0 { self.read(n) } else { 0 };
        (1u64 << n) - 1 + suffix
    }
}

/// Walk a codestream and return its marker segments in order.
///
/// Stops at EOC or at the first unrecognised marker. Each segment's payload length is
/// the exp-Golomb k0 value following the marker, after which the reader is byte-aligned.
pub fn parse_container(data: &[u8]) -> Vec<Segment> {
    let mut out = Vec::new();
    let mut pos = 0usize; // byte offset of the next marker
    while pos + 1 < data.len() {
        let mk = match Marker::from_u16(((data[pos] as u16) << 8) | data[pos + 1] as u16) {
            Some(m) => m,
            None => break,
        };
        if !mk.has_payload() {
            out.push(Segment { marker: mk, start: pos + 2, len: 0 });
            if mk == Marker::Eoc { break; }
            pos += 2;
            continue;
        }
        let mut br = BitReader::at(data, (pos + 2) * 8);
        let len = br.read_expgolomb_k0() as usize;
        br.align_to_byte();
        let start = br.bit_pos() / 8;
        if start + len > data.len() { break; }
        out.push(Segment { marker: mk, start, len });
        pos = start + len;
    }
    out
}

/// Find the first segment with `marker`.
pub fn find_segment(segs: &[Segment], marker: Marker) -> Option<Segment> {
    segs.iter().copied().find(|s| s.marker == marker)
}

/// Per-component (luma / chroma) parameters from the picture header.
#[derive(Clone, Debug, Default)]
pub struct ComponentHeader {
    pub num_threads_r: u32,
    pub num_chs: u32,
    pub use_cube_flags: bool,
    pub rvs_enable: bool,
    pub grfs_enable: bool,
    /// GRFS channel-selection flags, `num_chs` of them (empty if `grfs_enable` is false).
    pub grfs_channel_flags: Vec<bool>,
    /// `synthesis_tile_enable[c]` — the synthesis tile manager's flag, coded in the
    /// picture header right after this component's GRFS channel flags. (The
    /// component's LSBS flag is not here: LSBS codes it into the tool header; see
    /// [`parse_lsbs_flags`].)
    pub synthesis_tile_enable: bool,
}

/// The decoded JPEG-AI picture header (PIH substream).
#[derive(Clone, Debug, Default)]
pub struct PictureHeader {
    pub stream_profile_idc: u32,
    pub decoder_profile_id: u32,
    pub synthesis_transform_ids: Vec<u32>,
    pub level_idc: u32,
    pub img_width: u32,
    pub img_height: u32,
    pub diff_display_width: u32,
    pub diff_display_height: u32,
    pub bit_depth_idc: u32,
    /// Luma/chroma spatial subsampling factors (1 = 4:4:4).
    pub s_ver: u32,
    pub s_hor: u32,
    pub c_ver: u32,
    pub c_hor: u32,
    /// 0 = none, 1 = built-in YUV↔RGB, 2 = signalled 3×3 matrix (+ offset).
    pub colour_transform_idx: u32,
    pub colour_transform_matrix: Option<[u32; 9]>,
    pub colour_transform_offset: Option<[u32; 3]>,
    pub model_id: u32,
    pub num_threads_z: u32,
    /// `beta_displacement_log` (already de-biased by 2048).
    pub beta_displacement_log: [i32; 2],
    pub independent_beta_uv: bool,
    pub region_partitioning: bool,
    /// Two components: `[0]` = luma, `[1]` = chroma.
    pub components: [ComponentHeader; 2],
    pub gain_3d_enable: bool,
    /// Bit position just past the last field (before byte-alignment padding).
    pub end_bit: usize,
}

/// Parse the picture header (PIH payload). `start_bit` is the payload bit offset
/// (`Segment::start_bit`). A faithful port of the reference decode order:
/// prefix → `CCS_SGMM::decode_header` (`multi_threading_z`, quantizer β F-header,
/// `region_partitioning`, quantizer β O-header) → the per-component tool groups
/// (`multi_threading_r`, `num_chs`, `use_cube_flags`, `rvs_enable`, `grfs_enable` +
/// channel flags, `synthesis_tile_enable`) → `gain_3D_enable`.
///
/// Header-part placement follows the reference: the synthesis tile manager and
/// `gain_3D` are declared with `stream_header_part='pic_header'`, while LSBS keeps
/// the default `'tool_header'`, so `lsbs_enable_flag[0..1]` are the first two bits of
/// TON, not PIH fields ([`parse_lsbs_flags`]). The reference header trace records the
/// same order (`synthesis_tile_enable[0]` before `multi_threading_r[1]`,
/// `gain_3D_enable_flag` last) and `lsbs_enable_flag = 1` for both components.
///
/// `c_ver`/`c_hor` are coded only when the matching `s_ver`/`s_hor` is 1; otherwise
/// they are 2 and take no bits (reference `CodingEngine` header decode).
///
/// 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, `s_ver`/`s_hor` = 2) follow the reference syntax and are
/// decoded but untested against a real codestream.
// ---- 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 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
}

/// `lsbs_enable_flag[0..1]` (luma, chroma): the first two bits of the tool header
/// (TON), ahead of the post-filter chain that [`parse_tools_header`] reads from
/// `ton_start_bit + 2`. LSBS codes its enable flag into the tool header (the
/// reference's default header part), one flag per component.
pub fn parse_lsbs_flags(data: &[u8], ton_start_bit: usize) -> [bool; 2] {
    let mut br = BitReader::at(data, ton_start_bit);
    [br.read(1) == 1, br.read(1) == 1]
}

pub fn parse_picture_header(data: &[u8], start_bit: usize) -> PictureHeader {
    let mut br = BitReader::at(data, start_bit);
    let mut h = PictureHeader::default();

    h.stream_profile_idc = br.read(4) as u32;
    h.decoder_profile_id = br.read(4) as u32;
    let num_transforms = br.read(4) as u32 + 1;
    h.synthesis_transform_ids = (0..num_transforms).map(|_| br.read(4) as u32).collect();
    h.level_idc = br.read(8) as u32;
    h.img_width = br.read(16) as u32 + 64;
    h.img_height = br.read(16) as u32 + 64;
    h.diff_display_width = br.read(6) as u32;
    h.diff_display_height = br.read(6) as u32;
    h.bit_depth_idc = br.read(3) as u32;
    h.s_ver = br.read(1) as u32 + 1;
    h.s_hor = br.read(1) as u32 + 1;
    h.c_ver = if h.s_ver == 1 { br.read(1) as u32 + 1 } else { 2 };
    h.c_hor = if h.s_hor == 1 { br.read(1) as u32 + 1 } else { 2 };
    h.colour_transform_idx = br.read(2) as u32;
    if h.colour_transform_idx == 2 {
        let mut m = [0u32; 9];
        for v in m.iter_mut() { *v = br.read(8) as u32; }
        let mut o = [0u32; 3];
        for v in o.iter_mut() { *v = br.read(8) as u32; }
        h.colour_transform_matrix = Some(m);
        h.colour_transform_offset = Some(o);
    }
    h.model_id = br.read(4) as u32;

    // CCS_SGMM::decode_header
    h.num_threads_z = if br.read(1) == 1 { 1 << (br.read(2) + 1) } else { 1 };
    h.beta_displacement_log[0] = br.read(12) as i32 - 2048; // quantizer β F-header
    h.region_partitioning = br.read(1) == 1;
    if h.region_partitioning {
        let _num_ver = br.read(7) + 1;
        let _num_hor = br.read(7) + 1;
        if br.read(1) == 0 {
            let _hyper_overlap = br.read(2);
            let _mcm_overlap = br.read(4);
        }
    }
    h.independent_beta_uv = br.read(1) == 1; // quantizer β O-header
    h.beta_displacement_log[1] = if h.independent_beta_uv {
        br.read(12) as i32 - 2048
    } else {
        h.beta_displacement_log[0]
    };

    // per-component tool groups (0 = luma, 1 = chroma)
    for c in 0..2 {
        let comp = &mut h.components[c];
        comp.num_threads_r = if br.read(1) == 1 { 1 << (br.read(2) + 1) } else { 1 };
        comp.num_chs = br.read(8) as u32;
        // Only THREE flags precede the channel flags; the fourth PIH flag TRAILS
        // them. `res_var_scale`'s pair (rvs_enable, grfs_enable) sits immediately
        // before the channel flags it gates, and the synthesis tile manager's flag
        // follows the group. Reading four up-front shifts `grfs_channel_flags` one
        // bit late and corrupts them (while still landing on the right end_bit,
        // since the extra leading bit is cancelled by one fewer channel flag).
        comp.use_cube_flags = br.read(1) == 1;
        comp.rvs_enable = br.read(1) == 1;
        comp.grfs_enable = br.read(1) == 1;
        if comp.grfs_enable {
            comp.grfs_channel_flags = (0..comp.num_chs).map(|_| br.read(1) == 1).collect();
        }
        comp.synthesis_tile_enable = br.read(1) == 1;
    }
    h.gain_3d_enable = br.read(1) == 1;
    h.end_bit = br.bit_pos();
    h
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn bit_reader_is_msb_first() {
        // 0b1011_0010, 0b0100_0000
        let d = [0xB2u8, 0x40];
        let mut br = BitReader::new(&d);
        assert_eq!(br.read(1), 1);
        assert_eq!(br.read(3), 0b011);
        assert_eq!(br.read(4), 0b0010);
        assert_eq!(br.bit_pos(), 8);
        assert_eq!(br.read(2), 0b01, "continues MSB-first into the next byte");
    }

    #[test]
    fn read_past_end_returns_available_bits() {
        let d = [0xFFu8];
        let mut br = BitReader::at(&d, 4);
        // only 4 bits remain; the reference stops at the buffer end
        assert_eq!(br.read(8), 0b1111);
    }

    #[test]
    fn expgolomb_k0_round_values() {
        // k0: '1'->0 ; '010'->1 ; '011'->2 ; '00100'->3 ; '00111'->6
        let cases: [(&[u8], u64, u32); 3] = [
            (&[0b1000_0000], 0, 1),
            (&[0b0100_0000], 1, 3),
            (&[0b0011_1000], 6, 5),
        ];
        for (bytes, want, bits) in cases {
            let mut br = BitReader::new(bytes);
            assert_eq!(br.read_expgolomb_k0(), want, "value");
            assert_eq!(br.bit_pos(), bits as usize, "consumed bits");
        }
    }

    #[test]
    fn align_to_byte_is_idempotent_on_boundary() {
        let d = [0u8; 4];
        let mut br = BitReader::at(&d, 8);
        br.align_to_byte();
        assert_eq!(br.bit_pos(), 8, "already aligned");
        br.seek_bit(9);
        br.align_to_byte();
        assert_eq!(br.bit_pos(), 16);
    }

    /// MSB-first bit writer for synthetic headers.
    fn pack(fields: &[(u64, u32)]) -> Vec<u8> {
        let mut out = Vec::new();
        let mut n = 0usize;
        for &(v, w) in fields {
            for i in (0..w).rev() {
                if n % 8 == 0 { out.push(0); }
                let bit = ((v >> i) & 1) as u8;
                *out.last_mut().unwrap() |= bit << (7 - n % 8);
                n += 1;
            }
        }
        out.extend([0u8; 8]);
        out
    }

    /// `c_ver`/`c_hor` take no bits when `s_ver`/`s_hor` is 2 (reference
    /// `CodingEngine` header decode): the fields after them must not shift.
    #[test]
    fn chroma_factors_are_implicit_when_secondary_is_subsampled() {
        let prefix = [(0, 4), (2, 4), (0, 4), (2, 4), (52, 8), (192, 16), (192, 16), (0, 6), (0, 6), (0, 3)];
        // s_ver = 2, s_hor = 2 ⇒ no c_ver / c_hor bits; then colour_transform_idx = 1, model_id = 5
        let mut f: Vec<(u64, u32)> = prefix.to_vec();
        f.extend([(1, 1), (1, 1), (1, 2), (5, 4)]);
        let h = parse_picture_header(&pack(&f), 0);
        assert_eq!((h.s_ver, h.s_hor, h.c_ver, h.c_hor), (2, 2, 2, 2));
        assert_eq!((h.colour_transform_idx, h.model_id), (1, 5), "fields after s/c stay aligned");

        // s_ver = s_hor = 1 ⇒ c_ver / c_hor are coded (here 2 and 1)
        let mut f: Vec<(u64, u32)> = prefix.to_vec();
        f.extend([(0, 1), (0, 1), (1, 1), (0, 1), (1, 2), (5, 4)]);
        let h = parse_picture_header(&pack(&f), 0);
        assert_eq!((h.s_ver, h.s_hor, h.c_ver, h.c_hor), (1, 1, 2, 1));
        assert_eq!((h.colour_transform_idx, h.model_id), (1, 5));
    }

    #[test]
    fn parse_container_walks_markers_and_lengths() {
        // SOC | PIH(len 2, payload "ab") | EOC
        // expgolomb_k0(2) = '011' -> byte 0b0110_0000 after align
        let d = [0xFF, 0x80, 0xFF, 0x82, 0b0110_0000, b'a', b'b', 0xFF, 0x81];
        let segs = parse_container(&d);
        assert_eq!(segs.len(), 3, "{segs:?}");
        assert_eq!(segs[0].marker, Marker::Soc);
        assert_eq!(segs[1].marker, Marker::Pih);
        assert_eq!((segs[1].start, segs[1].len), (5, 2));
        assert_eq!(&d[segs[1].start..segs[1].start + segs[1].len], b"ab");
        assert_eq!(segs[2].marker, Marker::Eoc);
        assert!(find_segment(&segs, Marker::Ton).is_none());
    }
}