From 8fa22876f98972d730a7b8653a2432e4f1624fca Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 16 Jul 2026 04:39:36 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20jpegai=5Fbitstream=20=E2=80=94=20code?= =?UTF-8?q?stream=20container=20+=20MSB-first=20header=20bit=20reader?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The decoder's entry point. A codestream is a sequence of marker segments; each (except the marker-only SOC/EOC) is `marker(2B big-endian) + expgolomb_k0(size) + byte-align + payload`. The header substreams (PIH/TON) are plain bit-packed and read MSB-first — the headers are NOT arithmetic-coded (only the latent substreams are, and those go to jpegai_mans). BitReader faithful port of the reference EcLibDirect::read_bits — MSB-first within each byte, accumulated MSB-first, bit-granular cursor, reads past end return the available bits (matching the bounds check) read_expgolomb_k0 unsigned exp-Golomb k=0, for the segment payload lengths parse_container walk the markers -> [Segment{marker, start, len}] Validated against the real test256.bin codestream: full segment layout SOC/PIH(49)/TON(63)/RDI(1)/SORP(97)/SORS(47)/SOZ(83)/EOC, and the picture-header prefix decodes to 256x256, 4:4:4, colour_transform_idx=1, level 52. Cross-check: the SORP/SORS/SOZ payload sizes (97/47/83) equal the `mem` lengths the byte-exact mANS decoder consumes — the container parse and the entropy decode agree on the same bytes. Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_bitstream.rs | 204 ++++++++++++++++++ wai/wai-rs/src/lib.rs | 5 + .../tests/jpegai_bitstream_calibration.rs | 85 ++++++++ 3 files changed, 294 insertions(+) create mode 100644 wai/wai-rs/src/jpegai_bitstream.rs create mode 100644 wai/wai-rs/tests/jpegai_bitstream_calibration.rs diff --git a/wai/wai-rs/src/jpegai_bitstream.rs b/wai/wai-rs/src/jpegai_bitstream.rs new file mode 100644 index 000000000..741955af5 --- /dev/null +++ b/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 { + 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 { + 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 { + 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()); + } +} diff --git a/wai/wai-rs/src/lib.rs b/wai/wai-rs/src/lib.rs index b3897be85..dbf5e4003 100644 --- a/wai/wai-rs/src/lib.rs +++ b/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 diff --git a/wai/wai-rs/tests/jpegai_bitstream_calibration.rs b/wai/wai-rs/tests/jpegai_bitstream_calibration.rs new file mode 100644 index 000000000..1ce4f8bd7 --- /dev/null +++ b/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> { + 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 = 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 = (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"); +} -- 2.54.0 (Apple Git-157)