Skip to main content

rusty_h264_decoder/
cabac.rs

1//! CABAC arithmetic decoding engine (spec §9.3.3.2) + context initialization
2//! (§9.3.1.1). The literal-spec engine (codIRange/codIOffset, RenormD), which is
3//! bit-exact to openh264's optimized variant. Tables in [`crate::cabac_tables`].
4
5#[allow(unused_imports)]
6use alloc::borrow::ToOwned;
7#[allow(unused_imports)]
8use alloc::boxed::Box;
9#[allow(unused_imports)]
10use alloc::format;
11#[allow(unused_imports)]
12use alloc::string::{String, ToString};
13#[allow(unused_imports)]
14use alloc::vec;
15#[allow(unused_imports)]
16use alloc::vec::Vec;
17use rusty_h264_common::cabac_tables::{CTX_INIT, RANGE_LPS, STATE_TRANS};
18
19/// Profile-only bin census: how many bins of each class the engine decodes.
20/// The entropy stage's time divided by these counts gives ns/bin — the number
21/// that decides whether the engine or the syntax around it is the target.
22#[cfg(feature = "profile")]
23pub mod bin_census {
24    use core::sync::atomic::Ordering::Relaxed;
25    use rusty_h264_common::atomic::AtomicU64;
26    pub static DECISIONS: AtomicU64 = AtomicU64::new(0);
27    pub static BYPASSES: AtomicU64 = AtomicU64::new(0);
28    pub static TERMINATES: AtomicU64 = AtomicU64::new(0);
29    /// Decision bins whose renormalization shift was nonzero.
30    pub static RENORMS: AtomicU64 = AtomicU64::new(0);
31    pub fn reset() {
32        DECISIONS.store(0, Relaxed);
33        BYPASSES.store(0, Relaxed);
34        TERMINATES.store(0, Relaxed);
35    }
36    pub fn snapshot() -> (u64, u64, u64) {
37        (
38            DECISIONS.load(Relaxed),
39            BYPASSES.load(Relaxed),
40            TERMINATES.load(Relaxed),
41        )
42    }
43    pub fn renorms() -> u64 {
44        RENORMS.load(Relaxed)
45    }
46}
47
48/// FUSED per-(quartile, packed-state) record: `lps | trans_mps<<8 | trans_lps<<16`.
49///
50/// A context model is ONE byte: `state * 2 + mps` (0..=127) — ffmpeg/openh264's
51/// packing (H-35). The literal two-field form cost two loads and two stores per
52/// bin plus `1 - mps` arithmetic; packed, a bin is one byte load, one table
53/// lookup, one byte store, and `s & 1` for the value. This table folds the state
54/// transition AND the state-0 MPS flip (spec §9.3.3.2.1.1) into the lookup, so
55/// the decoded bins are identical by construction. Built at compile time from
56/// the spec tables: no init cost, no `OnceLock` check on the hot path.
57///
58/// Why: the serial chain of a decision bin ended with a LATE load — the
59/// transition table's address needs the LPS/MPS MASK, which exists only after
60/// the compare, so the context write-back (and every same-context successor
61/// bin: all unary and level-prefix loops re-read the context they just wrote)
62/// waited on a ~5-cycle L1 load issued at the chain's end. Folding both
63/// transition bytes into the SAME u32 the LPS quantity comes from makes them
64/// arrive EARLY (with the lps load, whose address needs only `s` and `q`),
65/// and the post-compare step becomes a 1-cycle shift-select:
66/// `(entry >> (8 + (mask & 8))) & 0xFF`. 2 KB, L1-resident like the tables it
67/// replaces on this path.
68const fn build_fused() -> [u32; 4 * 256] {
69    // Row stride 256, not 128: the packed context byte indexes the row DIRECTLY
70    // (`(q << 8) | s`, s < 256 by type), so the per-bin `& 127` mask that the
71    // 128-stride layout needed as its bounds proof is gone. Only s < 128 is
72    // ever populated or read (the packed-state invariant); the unused half of
73    // each row is never touched, so the L1 footprint is unchanged (2 KB live).
74    let mut t = [0u32; 4 * 256];
75    let mut q = 0;
76    while q < 4 {
77        let mut s = 0;
78        while s < 128 {
79            let lps = RANGE_LPS[s >> 1][q] as u32;
80            let tm = {
81                let mps = s as u8 & 1;
82                ((STATE_TRANS[s >> 1][1] << 1) | mps) as u32
83            };
84            let tl = {
85                let mps = s as u8 & 1;
86                let new_mps = if s >> 1 == 0 { 1 - mps } else { mps };
87                ((STATE_TRANS[s >> 1][0] << 1) | new_mps) as u32
88            };
89            t[q * 256 + s] = lps | (tm << 8) | (tl << 16);
90            s += 1;
91        }
92        q += 1;
93    }
94    t
95}
96static FUSED: [u32; 4 * 256] = build_fused();
97
98/// Bit position of the arithmetic offset field inside [`Engine::low`].
99const OFF: u32 = 41;
100/// Refill when fewer than this many buffered bits remain. 8 covers the worst
101/// single renormalization (6 bits) with margin; a 4-byte refill then lasts
102/// ~30 typical bins.
103const REFILL_AT: i32 = 8;
104
105#[cfg(feature = "cabac-trace")]
106mod trace {
107    use core::sync::atomic::{AtomicBool, AtomicU64, Ordering::Relaxed};
108    pub static ON: AtomicBool = AtomicBool::new(false);
109    pub static SYM: AtomicU64 = AtomicU64::new(0);
110    #[inline(never)]
111    pub fn emit(kind: &str, range: u32, offset: u64) {
112        let n = SYM.fetch_add(1, Relaxed);
113        eprintln!("{} {} r={} o={}", n, kind, range, offset);
114    }
115}
116
117/// Bring-up symbol trace hook. Compiled OUT unless the `cabac-trace` feature is
118/// on: the `if self.trace` test was a load + branch on EVERY bin (~154M per
119/// 1080p clip) and its cold `eprintln!` arm was inlined at every decode site --
120/// 54 print arms across the residual parser and the slice loop -- and the
121/// escaping call kept the engine state pinned to memory.
122#[inline(always)]
123fn tr(kind: &str, eng: &Engine) {
124    #[cfg(feature = "cabac-trace")]
125    if trace::ON.load(core::sync::atomic::Ordering::Relaxed) {
126        trace::emit(kind, eng.range, eng.low >> OFF);
127    }
128    #[cfg(not(feature = "cabac-trace"))]
129    let _ = (kind, eng);
130}
131
132/// Zero-filled 4-byte window at the end of the buffer -- cold by construction.
133#[cold]
134#[inline(never)]
135fn refill_tail(data: &[u8], p: usize) -> u32 {
136    let b = |i: usize| data.get(p + i).copied().unwrap_or(0) as u32;
137    (b(0) << 24) | (b(1) << 16) | (b(2) << 8) | b(3)
138}
139
140/// The arithmetic-decoder REGISTERS, by value. Four words. A parser pulls a copy
141/// out of the [`Cabac`] (`view`), decodes a whole block or syntax element against
142/// it -- LLVM keeps the four words in registers because nothing they alias is
143/// written -- and `commit`s once. Through `&mut Cabac` the same bins did 4 loads
144/// + 4 stores EACH: the context-byte store `ctx[idx]` lands in the same struct
145/// and LLVM will not prove it disjoint from `low`/`range`/`cnt`.
146#[derive(Clone, Copy, Debug)]
147pub struct Engine {
148    /// Next byte to load into the bit window.
149    byte_pos: usize,
150    /// FUSED offset+window register: `low = codIOffset * 2^41 + buf`, where
151    /// `buf < 2^41` holds the next `cnt` stream bits LEFT-ALIGNED at bit 40
152    /// downward, so renorm is ONE shift `low <<= n` and the next bits enter the
153    /// offset field by construction. Exact by invariant: `offset >= range <=>
154    /// low >= range << 41` (buffered bits can never flip the comparison); the
155    /// LPS subtraction cannot borrow into `buf`; `cnt <= 6 + 32 < 41`. Zero-fill
156    /// past the buffer end is preserved exactly (the fuzzer relies on it).
157    low: u64,
158    /// Valid buffered bits below the offset field.
159    cnt: i32,
160    range: u32,
161}
162
163/// The 460 context models, each packed as `state * 2 + mps`, in a 512-entry
164/// array so `ctx_idx & 511` is the index proof (one AND) instead of the
165/// `.min(459)` cmp+cmov that ran on every bin. Entries 460.. are never read.
166pub type Ctx = [u8; 512];
167
168impl Engine {
169    /// Appends 32 fresh stream bits directly below the current buffer fill
170    /// (zero-filled past the end of the data). Only called when `cnt <
171    /// REFILL_AT`, so the insert shift `9 - cnt` is always in `[2..=9]`.
172    ///
173    /// INLINED ALWAYS: as an out-of-line call inside every inlined bin decode it
174    /// was the one escape of the engine on the hot path. The past-the-end arm
175    /// stays cold.
176    #[inline(always)]
177    fn refill(&mut self, data: &[u8]) {
178        // REFUTED (2026-09-04): `if p + 4 <= len { data[p..=p+3] }` does NOT fold
179        // the four index checks -- `p + 4` may wrap in release, so LLVM cannot
180        // derive `p + 3 < len` from it -- and it added a panic path at every
181        // inlined refill site (+3,000 instrs in the slice loop). The range `get`
182        // is two compares and zero panic paths; keep it.
183        let v = match data.get(self.byte_pos..self.byte_pos + 4) {
184            Some(c) => u32::from_be_bytes([c[0], c[1], c[2], c[3]]),
185            None => refill_tail(data, self.byte_pos),
186        };
187        self.low |= (v as u64) << ((OFF as i32 - 32 - self.cnt) as u32);
188        self.byte_pos += 4;
189        self.cnt += 32;
190    }
191
192    /// Renormalization (spec 9.3.3.2.2): keep `range >= 256`. BRANCHLESS shift
193    /// count (`range <= 510` => `leading_zeros()-23` is exactly the iteration
194    /// count of the spec loop; 54-62% of decisions renormalize, a coin flip, so
195    /// a branch would mispredict); the refill is ONE shift of the fused register.
196    #[inline(always)]
197    fn renorm(&mut self, data: &[u8]) {
198        let n = self.range.leading_zeros() - 23;
199        self.range <<= n;
200        self.low <<= n;
201        self.cnt -= n as i32;
202        if self.cnt < REFILL_AT {
203            self.refill(data);
204        }
205    }
206
207    /// Decodes a context-coded bin (spec 9.3.3.2.1), updating the context model.
208    #[inline(always)]
209    pub fn decode_decision(&mut self, data: &[u8], ctx: &mut Ctx, ctx_idx: usize) -> u32 {
210        #[cfg(feature = "profile")]
211        bin_census::DECISIONS.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
212        tr("D", self);
213        // BRANCHLESS bin decode (H-35, the `get_cabac_inline` shape of ffmpeg).
214        // The LPS/MPS test is ~coin-flip on a well-adapted context, so a branch
215        // mispredicts constantly; derive an all-ones/zero MASK and select.
216        let ctx_idx = ctx_idx & 511;
217        let s = ctx[ctx_idx] as usize;
218        let q = ((self.range >> 6) & 3) as usize;
219        // ONE early load yields the LPS range AND both context transitions --
220        // see `build_fused` for why the transitions must not be a second,
221        // mask-addressed (late) load. `(q << 8) | s` is in-bounds by type.
222        let e = FUSED[(q << 8) | s];
223        let lps = e & 0xFF;
224        // PRECONDITION: `range >= 256` on entry, so `range - lps` (lps <= 240)
225        // stays positive and the sign test is a true "offset >= range" test.
226        debug_assert!(
227            self.range >= 256,
228            "renorm invariant broken: range={}",
229            self.range
230        );
231        self.range -= lps;
232        // mask = !0 when `offset >= range` (the LPS path), else 0 -- the sign
233        // trick in 64 bits against the SCALED range (values stay below 2^51).
234        let scaled = (self.range as u64) << OFF;
235        let mask64 = ((scaled as i64 - self.low as i64 - 1) >> 63) as u64;
236        let mask = mask64 as u32;
237        // LPS: offset -= range; range = lps.  MPS: both unchanged.
238        self.low -= scaled & mask64;
239        self.range = self.range.wrapping_add(lps.wrapping_sub(self.range) & mask);
240        // Both transitions arrived with the lps load; pick by mask with a
241        // shift (mask & 8 = 8 exactly on the LPS path).
242        ctx[ctx_idx] = ((e >> (8 + (mask & 8))) & 0xFF) as u8;
243        // MPS -> s&1; LPS -> (s&1)^1.
244        let bin = (s as u32 ^ mask) & 1;
245        #[cfg(feature = "profile")]
246        if self.range < 256 {
247            bin_census::RENORMS.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
248        }
249        self.renorm(data);
250        bin
251    }
252
253    /// Decodes a bypass (equiprobable) bin (spec 9.3.3.2.3).
254    #[inline(always)]
255    pub fn decode_bypass(&mut self, data: &[u8]) -> u32 {
256        #[cfg(feature = "profile")]
257        bin_census::BYPASSES.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
258        tr("B", self);
259        self.low <<= 1;
260        self.cnt -= 1;
261        if self.cnt < REFILL_AT {
262            self.refill(data);
263        }
264        // BRANCHLESS: a bypass bin is a coin flip by definition (signs, EG
265        // suffix bits), so the `if low >= scaled` branch mispredicted ~half the
266        // time. Same sign-mask select the decision path uses.
267        let scaled = (self.range as u64) << OFF;
268        let mask = ((scaled as i64 - self.low as i64 - 1) >> 63) as u64;
269        self.low -= scaled & mask;
270        (mask & 1) as u32
271    }
272
273    /// A bypass bin that FOLLOWS another bin of this engine with no refill in
274    /// between -- the coefficient SIGN. Every bin leaves `cnt >= REFILL_AT`
275    /// (renorm and bypass refill AFTER consuming), so one more bit cannot
276    /// exhaust the buffer, and whatever comes next is exact: a decision bin
277    /// renormalizes by at most 7 with cnt >= 7 and then refills; terminate
278    /// compares the offset field only. Saves the refill test per sign.
279    #[inline(always)]
280    pub fn decode_bypass_sign(&mut self) -> u32 {
281        #[cfg(feature = "profile")]
282        bin_census::BYPASSES.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
283        tr("B", self);
284        debug_assert!(self.cnt >= 1, "sign bypass needs a buffered bit");
285        self.low <<= 1;
286        self.cnt -= 1;
287        let scaled = (self.range as u64) << OFF;
288        let mask = ((scaled as i64 - self.low as i64 - 1) >> 63) as u64;
289        self.low -= scaled & mask;
290        (mask & 1) as u32
291    }
292
293    /// Decodes the terminate bin (spec 9.3.3.2.4); `true` ends the slice (or
294    /// marks I_PCM). No renormalization on terminate.
295    #[inline(always)]
296    pub fn decode_terminate(&mut self, data: &[u8]) -> bool {
297        #[cfg(feature = "profile")]
298        bin_census::TERMINATES.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
299        tr("T", self);
300        self.range -= 2;
301        if self.low >= (self.range as u64) << OFF {
302            true
303        } else {
304            self.renorm(data);
305            false
306        }
307    }
308}
309
310/// The CABAC decoder: arithmetic engine reading MSB-first from the RBSP plus the
311/// 460 adaptive context models.
312pub struct Cabac<'a> {
313    data: &'a [u8],
314    eng: Engine,
315    ctx: Ctx,
316}
317
318impl<'a> Cabac<'a> {
319    /// Initializes from the RBSP `data` at byte offset `start_byte` (the slice
320    /// data, byte-aligned past the header), the QP of the slice (clamped 0..51),
321    /// `cabac_init_idc`, and whether the slice is I/SI (spec 9.3.1).
322    pub fn new(data: &'a [u8], start_byte: usize, qp: i32, init_idc: u32, is_i: bool) -> Self {
323        let model = if is_i {
324            0
325        } else {
326            ((init_idc + 1) as usize).min(3)
327        };
328        let q = qp.clamp(0, 51);
329        let mut ctx = [0u8; 512];
330        for (i, c) in ctx.iter_mut().take(460).enumerate() {
331            let (m, n) = CTX_INIT[i][model];
332            let pre = (((m as i32 * q) >> 4) + n as i32).clamp(1, 126);
333            // Packed as state*2 + mps; same (state, mps) pair as the spec form.
334            *c = if pre <= 63 {
335                ((63 - pre) as u8) << 1
336            } else {
337                (((pre - 64) as u8) << 1) | 1
338            };
339        }
340        #[cfg(feature = "cabac-trace")]
341        trace::ON.store(
342            rusty_h264_common::knob("RH_CABAC_TRACE").is_some(),
343            core::sync::atomic::Ordering::Relaxed,
344        );
345        let mut e = Cabac {
346            data,
347            eng: Engine {
348                byte_pos: start_byte,
349                low: 0,
350                cnt: 0,
351                range: 510,
352            },
353            ctx,
354        };
355        e.eng.refill(data);
356        // codIOffset = first 9 bits: shift them from the buffer into the
357        // offset field -- the same fused move renorm makes every bin.
358        e.eng.low <<= 9;
359        e.eng.cnt -= 9;
360        e
361    }
362
363    /// Engine state `(codIRange, codIOffset)` -- for bring-up verification against
364    /// symbol 0 of the oracle (Brick 1.1). At slice start this is `(510, first-9-bits)`.
365    pub fn dbg_state(&self) -> (u32, u32) {
366        (self.eng.range, (self.eng.low >> OFF) as u32)
367    }
368
369    /// I_PCM sample position (spec 7.3.5 + 9.3.3.2.5). The PCM marker is a
370    /// terminate bin; after it decodes as 1, the flush output of the encoder is
371    /// already inside the borrowed offset bits of the engine, so the raw
372    /// `pcm_sample_*` bytes start at the consumed-bit position rounded up to
373    /// the next byte boundary (`pcm_alignment_zero_bit`s). `byte_pos*8 - cnt`
374    /// is that consumed position (offset-field bits count as read, buffered
375    /// bits do not). Valid only immediately after `decode_terminate()`
376    /// returned `true` (no renormalization has run since).
377    pub fn pcm_start_byte(&self) -> usize {
378        let consumed = self.eng.byte_pos as isize * 8 - self.eng.cnt as isize;
379        ((consumed + 7) >> 3) as usize
380    }
381
382    /// Re-initializes the arithmetic engine at absolute `byte` (spec 9.3.1.2,
383    /// invoked after the I_PCM samples), KEEPING the adaptive context models --
384    /// only the engine registers restart. Mirrors the tail of [`Cabac::new`].
385    pub fn reinit_at(&mut self, byte: usize) {
386        self.eng = Engine {
387            byte_pos: byte,
388            low: 0,
389            cnt: 0,
390            range: 510,
391        };
392        self.eng.refill(self.data);
393        self.eng.low <<= 9;
394        self.eng.cnt -= 9;
395    }
396
397    /// Split into `(data, engine copy, contexts)` for a register-resident run
398    /// over many bins; finish with [`Cabac::commit`].
399    #[inline(always)]
400    pub fn view(&mut self) -> (&'a [u8], Engine, &mut Ctx) {
401        (self.data, self.eng, &mut self.ctx)
402    }
403
404    /// Adopt the engine registers a [`Cabac::view`] run ended with.
405    #[inline(always)]
406    pub fn commit(&mut self, eng: Engine) {
407        self.eng = eng;
408    }
409
410    /// One context-coded bin. Single-bin convenience over [`Engine::decode_decision`];
411    /// multi-bin parsers take a [`Cabac::view`] instead.
412    #[inline(always)]
413    pub fn decode_decision(&mut self, ctx_idx: usize) -> u32 {
414        let mut e = self.eng;
415        let b = e.decode_decision(self.data, &mut self.ctx, ctx_idx);
416        self.eng = e;
417        b
418    }
419
420    /// One bypass bin (see [`Engine::decode_bypass`]).
421    #[inline(always)]
422    pub fn decode_bypass(&mut self) -> u32 {
423        let mut e = self.eng;
424        let b = e.decode_bypass(self.data);
425        self.eng = e;
426        b
427    }
428
429    /// The terminate bin (see [`Engine::decode_terminate`]).
430    #[inline(always)]
431    pub fn decode_terminate(&mut self) -> bool {
432        let mut e = self.eng;
433        let b = e.decode_terminate(self.data);
434        self.eng = e;
435        b
436    }
437
438    // NB: the byte offset where byte-aligned `pcm_sample` data resumes after an
439    // I_PCM under CABAC IS wired: `pcm_start_byte()` + `reinit_at()` above are
440    // the byte-realign/re-init pair, dispatched from the I_PCM arm of the
441    // decoder in mb16 and gated by tests/ipcm_cabac.rs against ffmpeg.
442}
443
444#[cfg(test)]
445mod tests {
446    use super::*;
447
448    /// Packed-form twin tables — the two-table formulation `FUSED` folded away.
449    /// Kept HERE (test-only) as the oracle that pins each `FUSED` field to the
450    /// spec tables; production reads only `FUSED`.
451    const fn build_lps_range() -> [u8; 4 * 128] {
452        let mut t = [0u8; 4 * 128];
453        let mut q = 0;
454        while q < 4 {
455            let mut s = 0;
456            while s < 128 {
457                t[q * 128 + s] = RANGE_LPS[s >> 1][q];
458                s += 1;
459            }
460            q += 1;
461        }
462        t
463    }
464    /// ONE transition table covering both paths: `[0..128)` MPS (state advances,
465    /// MPS unchanged), `[128..256)` LPS (state falls back; at state 0 the MPS
466    /// FLIPS per spec §9.3.3.2.1.1 — baked in, never branched).
467    const fn build_trans() -> [u8; 256] {
468        let mut t = [0u8; 256];
469        let mut s = 0;
470        while s < 128 {
471            let mps = s as u8 & 1;
472            t[s] = (STATE_TRANS[s >> 1][1] << 1) | mps;
473            let new_mps = if s >> 1 == 0 { 1 - mps } else { mps };
474            t[128 + s] = (STATE_TRANS[s >> 1][0] << 1) | new_mps;
475            s += 1;
476        }
477        t
478    }
479    static LPS_RANGE: [u8; 4 * 128] = build_lps_range();
480    static TRANS: [u8; 256] = build_trans();
481
482    /// Literal-spec CABAC *encoder* (§9.3.4), the inverse of [`Cabac`]. Used only
483    /// to validate the decoder by round-trip — encode a bin sequence, decode it,
484    /// assert equality. Encoder and decoder are independent algorithms (encode
485    /// vs decode), so a shared latent bug is implausible; a clean round-trip over
486    /// thousands of mixed bins exercises the full range/offset evolution, every
487    /// `RANGE_LPS`/`STATE_TRANS` entry reached, and the bypass/terminate paths.
488    struct Enc {
489        low: u32,
490        range: u32,
491        outstanding: u32,
492        first: bool,
493        bits: Vec<u8>,
494        ctx: Vec<(u8, u8)>, // (state, mps)
495    }
496
497    fn init_ctx(qp: i32, init_idc: u32, is_i: bool) -> Vec<(u8, u8)> {
498        let model = if is_i {
499            0
500        } else {
501            ((init_idc + 1) as usize).min(3)
502        };
503        let q = qp.clamp(0, 51);
504        (0..460)
505            .map(|i| {
506                let (m, n) = CTX_INIT[i][model];
507                let pre = (((m as i32 * q) >> 4) + n as i32).clamp(1, 126);
508                if pre <= 63 {
509                    ((63 - pre) as u8, 0)
510                } else {
511                    ((pre - 64) as u8, 1)
512                }
513            })
514            .collect()
515    }
516
517    impl Enc {
518        fn new(qp: i32, init_idc: u32, is_i: bool) -> Self {
519            Enc {
520                low: 0,
521                range: 510,
522                outstanding: 0,
523                first: true,
524                bits: Vec::new(),
525                ctx: init_ctx(qp, init_idc, is_i),
526            }
527        }
528
529        fn put_bit(&mut self, b: u32) {
530            if self.first {
531                self.first = false;
532            } else {
533                self.bits.push(b as u8);
534            }
535            while self.outstanding > 0 {
536                self.bits.push((1 - b) as u8);
537                self.outstanding -= 1;
538            }
539        }
540
541        /// RenormE (§9.3.4.3.3).
542        fn renorm(&mut self) {
543            while self.range < 256 {
544                if self.low < 256 {
545                    self.put_bit(0);
546                } else if self.low >= 512 {
547                    self.low -= 512;
548                    self.put_bit(1);
549                } else {
550                    self.low -= 256;
551                    self.outstanding += 1;
552                }
553                self.range <<= 1;
554                self.low <<= 1;
555            }
556        }
557
558        /// EncodeDecision (§9.3.4.3.1).
559        fn encode(&mut self, ctx_idx: usize, bin: u32) {
560            let (state, mps) = self.ctx[ctx_idx];
561            let q = ((self.range >> 6) & 3) as usize;
562            let lps = RANGE_LPS[state as usize][q] as u32;
563            self.range -= lps;
564            if bin != mps as u32 {
565                self.low += self.range;
566                self.range = lps;
567                let nm = if state == 0 { 1 - mps } else { mps };
568                self.ctx[ctx_idx] = (STATE_TRANS[state as usize][0], nm);
569            } else {
570                self.ctx[ctx_idx].0 = STATE_TRANS[state as usize][1];
571            }
572            self.renorm();
573        }
574
575        /// EncodeBypass (§9.3.4.3.2).
576        fn encode_bypass(&mut self, bin: u32) {
577            self.low <<= 1;
578            if bin != 0 {
579                self.low += self.range;
580            }
581            if self.low >= 1024 {
582                self.put_bit(1);
583                self.low -= 1024;
584            } else if self.low < 512 {
585                self.put_bit(0);
586            } else {
587                self.low -= 512;
588                self.outstanding += 1;
589            }
590        }
591
592        /// EncodeTerminate(1) + flush (§9.3.4.5 / EncodeFlush) — ends the stream.
593        fn finish(&mut self) -> Vec<u8> {
594            self.range -= 2;
595            self.low += self.range;
596            self.range = 2;
597            self.renorm();
598            self.put_bit((self.low >> 9) & 1);
599            let v = ((self.low >> 7) & 3) | 1;
600            self.bits.push(((v >> 1) & 1) as u8);
601            self.bits.push((v & 1) as u8);
602            // Pack MSB-first into bytes.
603            let mut out = vec![0u8; self.bits.len().div_ceil(8)];
604            for (i, &b) in self.bits.iter().enumerate() {
605                out[i / 8] |= b << (7 - (i % 8));
606            }
607            out
608        }
609    }
610
611    /// Deterministic xorshift RNG so the test is reproducible.
612    struct Rng(u32);
613    impl Rng {
614        fn next(&mut self) -> u32 {
615            self.0 ^= self.0 << 13;
616            self.0 ^= self.0 >> 17;
617            self.0 ^= self.0 << 5;
618            self.0
619        }
620    }
621
622    /// Encode a scripted mix of context-coded, bypass, and terminate bins, then
623    /// decode and assert every bin (and the terminate) round-trips exactly.
624    fn roundtrip(qp: i32, init_idc: u32, is_i: bool, seed: u32, n: usize) {
625        let mut rng = Rng(seed);
626        // (kind, ctx, bin): kind 0 = decision, 1 = bypass.
627        let mut script: Vec<(u8, usize, u32)> = Vec::with_capacity(n);
628        let mut enc = Enc::new(qp, init_idc, is_i);
629        for _ in 0..n {
630            let r = rng.next();
631            let kind = (r & 1) as u8;
632            let ctx = (r >> 1) as usize % 460;
633            let bin = (r >> 12) & 1;
634            script.push((kind, ctx, bin));
635            if kind == 0 {
636                enc.encode(ctx, bin);
637            } else {
638                enc.encode_bypass(bin);
639            }
640        }
641        let bytes = enc.finish();
642
643        let mut dec = Cabac::new(&bytes, 0, qp, init_idc, is_i);
644        for (i, &(kind, ctx, bin)) in script.iter().enumerate() {
645            let got = if kind == 0 {
646                dec.decode_decision(ctx)
647            } else {
648                dec.decode_bypass()
649            };
650            assert_eq!(got, bin, "bin {i} (kind {kind}, ctx {ctx}) mismatched");
651        }
652        assert!(
653            dec.decode_terminate(),
654            "terminate should signal end-of-stream"
655        );
656    }
657
658    #[test]
659    fn engine_roundtrip_many() {
660        // Sweep QP, init model, and many random scripts: every code path
661        // (LPS/MPS transitions across all 64 states, bypass, terminate, renorm).
662        for &qp in &[0, 12, 26, 37, 51] {
663            for &(idc, is_i) in &[(0u32, true), (0, false), (1, false), (2, false)] {
664                for seed in 1..=40u32 {
665                    roundtrip(
666                        qp,
667                        idc,
668                        is_i,
669                        seed.wrapping_mul(2654435761),
670                        seed as usize * 53,
671                    );
672                }
673            }
674        }
675    }
676
677    #[test]
678    fn engine_init_matches_spec() {
679        // ctxIdx 0 (I mb_type, m=20 n=-15) at QP 26: preCtxState =
680        // Clip3(1,126,(20*26>>4)-15) = 17 -> state 63-17 = 46, MPS 0.
681        let dec = Cabac::new(&[0xFF, 0xFF, 0xFF], 0, 26, 0, true);
682        // Packed as state*2 + mps (H-35): state 46, MPS 0 -> 92.
683        assert_eq!(dec.ctx[0] >> 1, 46, "state");
684        assert_eq!(dec.ctx[0] & 1, 0, "mps");
685        // Engine init: range 510, offset = first 9 bits of 0xFFFF = 0x1FF.
686        assert_eq!(dec.eng.range, 510);
687        assert_eq!(dec.dbg_state().1, 0x1FF);
688    }
689
690    /// H-35 oracle: for EVERY packed state and range quartile, the packed tables
691    /// must reproduce the literal spec derivation (RangeLPS, the bin value, and
692    /// both transitions including the state-0 MPS flip) exactly. 512 cases —
693    /// cheaper and stricter than trusting a corpus.
694    #[test]
695    fn packed_state_tables_match_spec_form() {
696        for s in 0usize..128 {
697            let (state, mps) = ((s >> 1) as u8, (s & 1) as u8);
698            for q in 0usize..4 {
699                assert_eq!(
700                    LPS_RANGE[q * 128 + s],
701                    RANGE_LPS[state as usize][q],
702                    "lps s={s} q={q}"
703                );
704            }
705            // MPS half: bin == mps, state advances, mps unchanged.
706            let mps_t = TRANS[s];
707            assert_eq!(
708                mps_t >> 1,
709                STATE_TRANS[state as usize][1],
710                "mps-trans state s={s}"
711            );
712            assert_eq!(mps_t & 1, mps, "mps-trans mps s={s}");
713            // LPS half: bin == 1-mps, state falls back, mps flips only at state 0.
714            let lps_t = TRANS[128 + s];
715            let want_mps = if state == 0 { 1 - mps } else { mps };
716            assert_eq!(
717                lps_t >> 1,
718                STATE_TRANS[state as usize][0],
719                "lps-trans state s={s}"
720            );
721            assert_eq!(lps_t & 1, want_mps, "lps-trans mps s={s}");
722        }
723    }
724
725    /// H-35 oracle #2: the BRANCHLESS mask arithmetic must equal the literal
726    /// `if offset >= range` form for every (range, offset, state) combination
727    /// the engine can present — the mask, the two conditional updates, the
728    /// transition-table half selection, and the bin value. This is the whole
729    /// risk surface of the branchless rewrite, checked exhaustively rather than
730    /// inferred from a corpus that happens to decode.
731    #[test]
732    fn branchless_mask_matches_conditional_form() {
733        // Reachable domain only: renorm guarantees `range` in 256..=510 on entry
734        // and the spec invariant `offset < range` holds throughout. (Widening
735        // past this tests states the engine cannot present — and the wrapped
736        // `range - lps` there makes BOTH forms meaningless, not just one.)
737        for s in 0usize..128 {
738            for range in [256u32, 257, 300, 383, 384, 400, 448, 509, 510] {
739                for offset in [0u32, 1, 127, 128, 255, 256, 300, 383, 384, 509] {
740                    if offset >= range {
741                        continue;
742                    }
743                    let q = ((range >> 6) & 3) as usize;
744                    let lps = LPS_RANGE[q * 128 + s] as u32;
745                    let r1 = range.wrapping_sub(lps);
746                    // literal spec form
747                    let (mut lr, mut lo, lbin, lctx) = if offset >= r1 {
748                        (lps, offset - r1, (s as u32 & 1) ^ 1, TRANS[128 + s])
749                    } else {
750                        (r1, offset, s as u32 & 1, TRANS[s])
751                    };
752                    // branchless form, exactly as `decode_decision` computes it
753                    let mask = ((r1 as i32 - offset as i32 - 1) >> 31) as u32;
754                    let bo = offset - (r1 & mask);
755                    let br = r1.wrapping_add(lps.wrapping_sub(r1) & mask);
756                    let bctx = TRANS[s | (mask as usize & 128)];
757                    let bbin = (s as u32 ^ mask) & 1;
758                    // (silence unused-mut on the literal bindings)
759                    lr += 0;
760                    lo += 0;
761                    assert_eq!(
762                        (lr, lo, lbin, lctx),
763                        (br, bo, bbin, bctx),
764                        "s={s} range={range} offset={offset}"
765                    );
766                }
767            }
768        }
769    }
770
771    #[test]
772    fn tables_match_spec_boundaries() {
773        assert_eq!(RANGE_LPS[0], [128, 176, 208, 240]);
774        assert_eq!(RANGE_LPS[63], [2, 2, 2, 2]);
775        assert_eq!(STATE_TRANS[0], [0, 1]);
776        assert_eq!(STATE_TRANS[63], [63, 63]);
777        assert_eq!(CTX_INIT[0][0], (20, -15));
778    }
779}