rusty_h264-encoder 0.3.0

//! CABAC arithmetic *encoder* engine (spec §9.3.4) — the exact forward inverse of
//! the decoder's [`rusty_h264_decoder::cabac`] engine. Promoted verbatim from the
//! decoder's round-trip-validated test encoder (`engine_roundtrip_many` sweeps
//! QP × init-model × 40 seeds), so the range/offset evolution, every
//! `RANGE_LPS`/`STATE_TRANS` transition, and the bypass/terminate/flush paths are
//! already proven bit-exact against [`Cabac`]. This module only wraps that engine
//! in a syntax-facing API (`encode_decision`/`encode_bypass`/`encode_terminate`)
//! and the 460 adaptive context models.
//!
//! Tables (`CTX_INIT`, `RANGE_LPS`, `STATE_TRANS`) are shared with the decoder via
//! [`rusty_h264_common::cabac_tables`] — one source of truth, no drift.

use rusty_h264_common::cabac_tables::{CTX_INIT, RANGE_LPS, STATE_TRANS};

/// Initialise the 460 context models `(state, mps)` from `CTX_INIT` (spec §9.3.1.1).
/// Identical to the decoder's `Cabac::new` context init.
fn init_ctx(qp: i32, init_idc: u32, is_i: bool) -> Vec<(u8, u8)> {
    let model = if is_i { 0 } else { ((init_idc + 1) as usize).min(3) };
    let q = qp.clamp(0, 51);
    (0..460)
        .map(|i| {
            let (m, n) = CTX_INIT[i][model];
            let pre = (((m as i32 * q) >> 4) + n as i32).clamp(1, 126);
            if pre <= 63 {
                ((63 - pre) as u8, 0)
            } else {
                ((pre - 64) as u8, 1)
            }
        })
        .collect()
}

/// The CABAC arithmetic encoder: the low/range interval coder (spec §9.3.4) plus
/// the 460 adaptive context models. `bits` accumulates the raw output bit stream
/// (delayed by `put_bit`'s carry/outstanding logic); `into_bytes` packs it
/// MSB-first once the final `encode_terminate(true)` has flushed the coder.
pub struct CabacEncoder {
    low: u32,
    range: u32,
    outstanding: u32,
    first: bool,
    bits: Vec<u8>,
    ctx: Vec<(u8, u8)>,
    /// Running count of bins emitted — the RD bit-cost proxy (each context/bypass
    /// bin is ~1 coded bit; adaptive contexts make it fractional, but the *count*
    /// is the cheap monotone cost surrogate the mode decision can use).
    pub bins: u64,
}

impl CabacEncoder {
    /// New encoder with contexts initialised for `qp` / `init_idc` / slice type.
    pub fn new(qp: i32, init_idc: u32, is_i: bool) -> Self {
        CabacEncoder {
            low: 0,
            range: 510,
            outstanding: 0,
            first: true,
            bits: Vec::new(),
            ctx: init_ctx(qp, init_idc, is_i),
            bins: 0,
        }
    }

    /// Emit a resolved bit plus any carry-delayed `outstanding` bits (spec's
    /// bit-with-carry PutBit).
    fn put_bit(&mut self, b: u32) {
        if self.first {
            self.first = false;
        } else {
            self.bits.push(b as u8);
        }
        while self.outstanding > 0 {
            self.bits.push((1 - b) as u8);
            self.outstanding -= 1;
        }
    }

    /// RenormE (§9.3.4.3.3).
    fn renorm(&mut self) {
        while self.range < 256 {
            if self.low < 256 {
                self.put_bit(0);
            } else if self.low >= 512 {
                self.low -= 512;
                self.put_bit(1);
            } else {
                self.low -= 256;
                self.outstanding += 1;
            }
            self.range <<= 1;
            self.low <<= 1;
        }
    }

    /// EncodeDecision (§9.3.4.3.1) — code one context-adaptive bin and update the model.
    pub fn encode_decision(&mut self, ctx_idx: usize, bin: u32) {
        let (state, mps) = self.ctx[ctx_idx];
        let q = ((self.range >> 6) & 3) as usize;
        let lps = RANGE_LPS[state as usize][q] as u32;
        self.range -= lps;
        if bin != mps as u32 {
            self.low += self.range;
            self.range = lps;
            let nm = if state == 0 { 1 - mps } else { mps };
            self.ctx[ctx_idx] = (STATE_TRANS[state as usize][0], nm);
        } else {
            self.ctx[ctx_idx].0 = STATE_TRANS[state as usize][1];
        }
        self.renorm();
        self.bins += 1;
    }

    /// EncodeBypass (§9.3.4.3.2) — code one equiprobable bin (no context).
    pub fn encode_bypass(&mut self, bin: u32) {
        self.low <<= 1;
        if bin != 0 {
            self.low += self.range;
        }
        if self.low >= 1024 {
            self.put_bit(1);
            self.low -= 1024;
        } else if self.low < 512 {
            self.put_bit(0);
        } else {
            self.low -= 512;
            self.outstanding += 1;
        }
        self.bins += 1;
    }

    /// Code `n` bypass bins of `val`, MSB first (unsigned bypass strings). Reserved
    /// for the P/B inter syntax (mvd / ref_idx bypass suffixes) — see CABAC-4.
    #[allow(dead_code)]
    pub fn encode_bypass_bits(&mut self, val: u32, n: u32) {
        for i in (0..n).rev() {
            self.encode_bypass((val >> i) & 1);
        }
    }

    /// EncodeTerminate (§9.3.4.5) — code the `end_of_slice_flag`. `end == false`
    /// between MBs (more to come); `end == true` on the last MB, which also runs
    /// EncodeFlush (§9.3.4.6) to close the stream. After `end == true` call
    /// [`into_bytes`](Self::into_bytes).
    pub fn encode_terminate(&mut self, end: bool) {
        self.range -= 2;
        if !end {
            self.renorm();
        } else {
            self.low += self.range;
            self.range = 2;
            self.renorm();
            // EncodeFlush bit output.
            self.put_bit((self.low >> 9) & 1);
            let v = ((self.low >> 7) & 3) | 1;
            self.bits.push(((v >> 1) & 1) as u8);
            self.bits.push((v & 1) as u8);
        }
        self.bins += 1;
    }

    /// Pack the accumulated bits MSB-first into bytes. Call after the final
    /// `encode_terminate(true)`. The output is byte-aligned (EncodeFlush guarantees
    /// the closing `1` + alignment), ready to append after the byte-aligned slice
    /// header.
    pub fn into_bytes(self) -> Vec<u8> {
        let mut out = vec![0u8; self.bits.len().div_ceil(8)];
        for (i, &b) in self.bits.iter().enumerate() {
            out[i / 8] |= b << (7 - (i % 8));
        }
        out
    }
}

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

    struct Rng(u32);
    impl Rng {
        fn next(&mut self) -> u32 {
            self.0 ^= self.0 << 13;
            self.0 ^= self.0 >> 17;
            self.0 ^= self.0 << 5;
            self.0
        }
    }

    /// The promoted engine must still round-trip through the decoder exactly — the
    /// same sweep the decoder validated its own test encoder with, now against this
    /// production module.
    fn roundtrip(qp: i32, init_idc: u32, is_i: bool, seed: u32, n: usize) {
        let mut rng = Rng(seed);
        let mut script: Vec<(u8, usize, u32)> = Vec::with_capacity(n);
        let mut enc = CabacEncoder::new(qp, init_idc, is_i);
        for _ in 0..n {
            let r = rng.next();
            let kind = (r & 1) as u8;
            let ctx = (r >> 1) as usize % 460;
            let bin = (r >> 12) & 1;
            script.push((kind, ctx, bin));
            if kind == 0 {
                enc.encode_decision(ctx, bin);
            } else {
                enc.encode_bypass(bin);
            }
        }
        enc.encode_terminate(true);
        let bytes = enc.into_bytes();

        let mut dec = Cabac::new(&bytes, 0, qp, init_idc, is_i);
        for (i, &(kind, ctx, bin)) in script.iter().enumerate() {
            let got = if kind == 0 {
                dec.decode_decision(ctx)
            } else {
                dec.decode_bypass()
            };
            assert_eq!(got, bin, "bin {i} (kind {kind}, ctx {ctx}) mismatched");
        }
        assert!(dec.decode_terminate(), "terminate should signal end-of-stream");
    }

    #[test]
    fn engine_roundtrip_many() {
        for &qp in &[0, 12, 26, 37, 51] {
            for &(idc, is_i) in &[(0u32, true), (0, false), (1, false), (2, false)] {
                for seed in 1..=40u32 {
                    roundtrip(qp, idc, is_i, seed.wrapping_mul(2654435761), seed as usize * 53);
                }
            }
        }
    }
}