wai/wai-rs/src/jpegai_bitstream.rs

@@ -0,0 +1,204 @@
+//! 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.
+
+/// 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
+    }
+
+    /// 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)
+}
+
+#[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);
+    }
+
+    #[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());
+    }
+}

wai/wai-rs/src/lib.rs

@@ -87,6 +87,11 @@ pub mod jpegai_pipeline;
 #[cfg(feature = "neural_int")]
 pub mod jpegai_icci;
 
+// JPEG-AI bitstream container (marker segments) + the MSB-first header bit reader —
+// the decoder's entry point.
+#[cfg(feature = "neural_int")]
+pub mod jpegai_bitstream;
+
 // Deterministic integer generative synthesis (`wai.generate.*`) — the
 // `intent=create` analog of the integer-exact codecs: a generator whose entire
 // pipeline is integer/fixed-point, pinned to (model, prompt, seed, steps), so the

wai/wai-rs/tests/jpegai_bitstream_calibration.rs

@@ -0,0 +1,85 @@
+//! Calibration for [`wai::jpegai_bitstream`] against a REAL JPEG-AI codestream
+//! (`test256.bin`): walk the marker segments and check the layout, then replay the
+//! picture header's leading fields with the bit reader and match the reference values.
+//!
+//! The strongest check here is a cross-validation: the SORP / SORS / SOZ payload sizes
+//! recovered by the container parse must equal the `mem` lengths in the entropy fixture
+//! that the byte-exact mANS decoder consumes — two independent routes to the same bytes.
+//!
+//! env: WAI_JPEGAI_CODESTREAM = path to test256.bin
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_bitstream::{find_segment, parse_container, BitReader, Marker};
+
+fn codestream() -> Option<Vec<u8>> {
+    match std::env::var("WAI_JPEGAI_CODESTREAM") {
+        Ok(p) if std::path::Path::new(&p).exists() => Some(std::fs::read(p).unwrap()),
+        _ => { eprintln!("skip: set WAI_JPEGAI_CODESTREAM to test256.bin"); None }
+    }
+}
+
+#[test]
+fn container_layout_matches_reference_codestream() {
+    let Some(data) = codestream() else { return };
+    let segs = parse_container(&data);
+    for s in &segs {
+        eprintln!("  {:?} payload bytes {}..{} (len {})", s.marker, s.start, s.start + s.len, s.len);
+    }
+    let order: Vec<Marker> = segs.iter().map(|s| s.marker).collect();
+    assert_eq!(
+        order,
+        vec![Marker::Soc, Marker::Pih, Marker::Ton, Marker::Rdi, Marker::Sorp, Marker::Sors, Marker::Soz, Marker::Eoc],
+        "codestream segment order"
+    );
+    let pih = find_segment(&segs, Marker::Pih).unwrap();
+    assert_eq!((pih.start, pih.len), (6, 49), "PIH payload");
+    let ton = find_segment(&segs, Marker::Ton).unwrap();
+    assert_eq!((ton.start, ton.len), (59, 63), "TON payload");
+
+    // Cross-validation: the entropy substream sizes must equal the mANS fixture `mem`
+    // lengths (SORP=97, SORS=47, SOZ=83) — the container parse and the byte-exact
+    // entropy decode agree on the same bytes from opposite directions.
+    assert_eq!(find_segment(&segs, Marker::Sorp).unwrap().len, 97, "SORP == mANS mem");
+    assert_eq!(find_segment(&segs, Marker::Sors).unwrap().len, 47, "SORS == mANS mem");
+    assert_eq!(find_segment(&segs, Marker::Soz).unwrap().len, 83, "SOZ == mANS mem");
+}
+
+#[test]
+fn picture_header_leading_fields_match_reference() {
+    let Some(data) = codestream() else { return };
+    let segs = parse_container(&data);
+    let pih = find_segment(&segs, Marker::Pih).expect("PIH present");
+    let mut br = BitReader::at(&data, pih.start_bit());
+    assert_eq!(br.bit_pos(), 48, "PIH payload starts at bit 48");
+
+    // The picture header's fixed prefix, per the reference syntax.
+    let stream_profile_idc = br.read(4);
+    let decoder_profile_id = br.read(4);
+    let num_synthesis_transforms_minus1 = br.read(4);
+    let synthesis_transform_id: Vec<u64> = (0..=num_synthesis_transforms_minus1).map(|_| br.read(4)).collect();
+    let level_idc = br.read(8);
+    let img_width = br.read(16) + 64;
+    let img_height = br.read(16) + 64;
+    let _diff_display_w = br.read(6);
+    let _diff_display_h = br.read(6);
+    let bit_depth_idc = br.read(3);
+    let (s_ver, s_hor) = (br.read(1) + 1, br.read(1) + 1);
+    let (c_ver, c_hor) = (br.read(1) + 1, br.read(1) + 1);
+    let colour_transform_idx = br.read(2);
+    let model_id = br.read(4);
+
+    eprintln!("profile={stream_profile_idc} decoder={decoder_profile_id} level={level_idc} \
+               {img_width}x{img_height} bd_idc={bit_depth_idc} colour_idx={colour_transform_idx} model={model_id}");
+
+    assert_eq!(stream_profile_idc, 0);
+    assert_eq!(decoder_profile_id, 2);
+    assert_eq!(num_synthesis_transforms_minus1, 2);
+    assert_eq!(synthesis_transform_id, vec![2, 1, 0]);
+    assert_eq!(level_idc, 52);
+    assert_eq!((img_width, img_height), (256, 256), "test256 is 256x256");
+    assert_eq!(bit_depth_idc, 0);
+    assert_eq!((s_ver, s_hor, c_ver, c_hor), (1, 1, 1, 1), "4:4:4");
+    assert_eq!(colour_transform_idx, 1, "built-in YUV<->RGB");
+    assert_eq!(model_id, 0);
+    assert_eq!(br.bit_pos(), 137, "cursor lands where the reference trace does");
+}