//! Rust port of CompressAI's rANS decoder (compressai/cpp_exts/rans/
//! rans_interface.cpp + third_party/ryg_rans/rans64.h).
//!
//! Byte-exact compatible with the upstream encoder — the Rust native
//! sink and the browser JS port (wai-web/demo/rans.js) both decode the
//! same bitstream to the same int32 symbols. precision=16, bypass=4.

const PRECISION:        u32 = 16;
const BYPASS_PRECISION: u32 = 4;
const MAX_BYPASS_VAL:   u32 = (1 << BYPASS_PRECISION) - 1;
const RANS64_L:         u64 = 1u64 << 31;

pub struct RansDecoder<'a> {
    words: &'a [u8],  // big-byte view; we read u32 LE words at idx
    idx:   usize,     // word index (i.e., 4*idx is byte offset)
    state: u64,
    /// Set when a renormalisation would have run past the end of the bitstream.
    /// A well-formed stream never sets it. A truncated one does, and the decoder
    /// then stops consuming rather than panicking the sink that fed it.
    exhausted: bool,
}

impl<'a> RansDecoder<'a> {
    pub fn new(bytes: &'a [u8]) -> Result<Self, &'static str> {
        if !bytes.len().is_multiple_of(4) {
            return Err("rANS bitstream length not /4");
        }
        if bytes.len() < 8 {
            return Err("rANS bitstream too short to seed 64-bit state");
        }
        let w0 = u32::from_le_bytes(bytes[0..4].try_into().unwrap()) as u64;
        let w1 = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as u64;
        Ok(Self { words: bytes, idx: 2, state: (w1 << 32) | w0, exhausted: false })
    }

    /// Whether the stream ran out during a renormalisation. Decoded values are
    /// not meaningful past this point; a conforming stream never sets it.
    pub fn is_exhausted(&self) -> bool {
        self.exhausted
    }

    /// Whether the stream has ended canonically: it never ran out, every word has been
    /// read, and the state is `2^31`, the state an encoder starts from. Called after a
    /// decode's last symbol. A stream that ends otherwise (an unread word after the
    /// last symbol, or a state the encoder did not start from) codes the same symbols
    /// as the canonical stream, so a decoder that requires this accepts one stream per
    /// sequence of symbols.
    pub fn is_canonical_end(&self) -> bool {
        !self.exhausted && self.state == RANS64_L && self.idx * 4 == self.words.len()
    }

    #[inline] fn renorm(&mut self) {
        if self.state < RANS64_L {
            let off = self.idx * 4;
            // A truncated or fully-consumed stream must not panic a sink that
            // was handed untrusted bytes. A conforming stream always has the
            // word this branch wants, so the guard is a no-op on every
            // well-formed decode.
            let Some(w) = self.words.get(off..off + 4) else {
                self.exhausted = true;
                return;
            };
            let next = u32::from_le_bytes(w.try_into().unwrap()) as u64;
            self.state = (self.state << 32) | next;
            self.idx += 1;
        }
    }

    fn get_bits(&mut self, n: u32) -> u32 {
        let mask = (1u64 << n) - 1;
        let val  = (self.state & mask) as u32;
        self.state >>= n;
        self.renorm();
        val
    }

    /// Decode one symbol against `cdf[0..cdf_len]` with per-channel `offset`.
    pub fn decode(&mut self, cdf: &[i32], cdf_len: usize, offset: i32) -> i32 {
        // A table shorter than two entries has no alphabet, and `cdf_len - 1`
        // below would underflow; a `cdf_len` past the slice would index out of
        // bounds. Both are caller-contract violations that a conforming codec
        // never produces — handled here so the function is total rather than
        // merely usually-safe.
        if cdf_len < 2 || cdf_len > cdf.len() {
            return offset;
        }
        let max_value = (cdf_len as i32) - 2;
        let cum_freq  = (self.state & ((1u64 << PRECISION) - 1)) as i32;
        // Binary search for the largest s such that cdf[s] <= cum_freq.
        let (mut lo, mut hi) = (1, cdf_len - 1);
        while lo < hi {
            let mid = (lo + hi) / 2;
            if cdf[mid] > cum_freq { hi = mid; } else { lo = mid + 1; }
        }
        let s = (lo - 1) as i32;
        let start = cdf[s as usize] as u64;
        let freq  = (cdf[(s + 1) as usize] - cdf[s as usize]) as u64;
        self.state = freq * (self.state >> PRECISION)
                   + (self.state & ((1u64 << PRECISION) - 1)) - start;
        self.renorm();

        let mut value = s;
        if value == max_value {
            // `raw_val` is 32 bits and each group carries BYPASS_PRECISION of
            // them, so at most this many groups can contribute anything. The
            // run length is read *from the stream*, so without a bound a
            // malformed bitstream drives the shift distance to 32 or beyond —
            // which panics in debug and masks in release, i.e. the same bytes
            // decode two different ways depending on how the sink was compiled.
            // A conforming encoder never emits more than MAX_BYPASS_GROUPS, so
            // every clamp here is a no-op on a well-formed stream.
            const MAX_BYPASS_GROUPS: u32 = 32 / BYPASS_PRECISION;
            let mut val = self.get_bits(BYPASS_PRECISION);
            let mut n_bypass = val;
            let mut groups = 0u32;
            while val == MAX_BYPASS_VAL && groups < MAX_BYPASS_GROUPS && !self.exhausted {
                val = self.get_bits(BYPASS_PRECISION);
                n_bypass = n_bypass.saturating_add(val);
                groups += 1;
            }
            let n_bypass = n_bypass.min(MAX_BYPASS_GROUPS);
            let mut raw_val: u32 = 0;
            for j in 0..n_bypass {
                let v = self.get_bits(BYPASS_PRECISION);
                raw_val |= v << (j * BYPASS_PRECISION);
            }
            value = (raw_val >> 1) as i32;
            if (raw_val & 1) != 0 { value = (-value).saturating_sub(1); }
            else                  { value = value.saturating_add(max_value); }
        }
        value.saturating_add(offset)
    }

    /// Decode one symbol of a **categorical** distribution (GLC / VQ codebook):
    /// every index `0..K-1` is in `cdf[0..=K]` (length `K+1`, `cdf[K] = 2^P`),
    /// no bypass, no offset. Same rANS state machine as [`decode`], minus the
    /// out-of-range escape. `cdf` is the integer table from
    /// `neural::glc::pmf_to_quantized_cdf`.
    pub fn decode_categorical(&mut self, cdf: &[u32]) -> u32 {
        let cum_freq = (self.state & ((1u64 << PRECISION) - 1)) as u32;
        let (mut lo, mut hi) = (1usize, cdf.len() - 1);
        while lo < hi {
            let mid = (lo + hi) / 2;
            if cdf[mid] > cum_freq { hi = mid; } else { lo = mid + 1; }
        }
        let s = lo - 1;
        let start = cdf[s] as u64;
        let freq = (cdf[s + 1] - cdf[s]) as u64;
        self.state = freq * (self.state >> PRECISION)
                   + (self.state & ((1u64 << PRECISION) - 1)) - start;
        self.renorm();
        s as u32
    }
}

// ─────────────────────────────────────────────────────────────────────
// Kani proof harnesses — the determinism claim, checked rather than sampled
// ─────────────────────────────────────────────────────────────────────
//
// WAI's contract is that an integer decode is byte-identical on every machine.
// The cross-architecture matrix samples that on four targets with well-formed
// streams; these harnesses reason over *all* inputs in a bounded window.
//
// The property that actually carries the claim is arithmetic totality. Rust
// integer overflow **panics in debug and wraps in release**, so an overflowing
// decoder is not one implementation — it is two, and which one a sink runs
// depends on how it was compiled. Proving the arithmetic cannot overflow is
// therefore proving the bitstream means one thing.
//
// The pair below is deliberate: the same decode, once without the rANS state
// invariant and once with it, so the proof isolates exactly which precondition
// the guarantee rests on rather than asserting safety unconditionally.

/// The rANS state invariant the encoder maintains: after seeding and after each
/// renormalisation the state is a 63-bit quantity. `RansDecoder::new` seeds from
/// caller-supplied bytes and does **not** establish it, which is what the paired
/// harnesses below make precise.
#[cfg(kani)]
const RANS64_STATE_BOUND: u64 = 1u64 << 63;

/// A well-formed 3-symbol categorical CDF: `cdf[0] = 0`, strictly increasing,
/// `cdf[K] = 2^PRECISION`. Symbolic in the interior points, so the harness
/// reasons over every admissible table of this size rather than one example.
#[cfg(kani)]
fn any_wellformed_cdf3() -> [u32; 4] {
    let a: u32 = kani::any();
    let b: u32 = kani::any();
    kani::assume(0 < a && a < b && b < (1u32 << PRECISION));
    [0, a, b, 1u32 << PRECISION]
}

/// **With the state invariant**: the decode is total (no overflow, no
/// underflow, no out-of-bounds) and the symbol it returns is inside the
/// alphabet, for every admissible CDF and every state in range.
///
/// This is the theorem the byte-exactness claim needs: on a stream the encoder
/// produced, the arithmetic cannot overflow, so debug and release agree and the
/// bitstream denotes exactly one symbol sequence.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(8)]
fn kani_decode_categorical_total_under_state_invariant() {
    let cdf = any_wellformed_cdf3();
    let bytes: [u8; 16] = kani::any();
    let mut d = match RansDecoder::new(&bytes) {
        Ok(d) => d,
        Err(_) => return,
    };
    // The invariant the encoder establishes and every renorm preserves.
    kani::assume(d.state < RANS64_STATE_BOUND);
    let s = d.decode_categorical(&cdf);
    kani::assert(
        (s as usize) < cdf.len() - 1,
        "decoded symbol lies inside the alphabet",
    );
    // Non-vacuity, checked rather than asserted. A harness whose assumptions
    // admit no state, or which returns before the call, would satisfy every
    // assertion trivially. This cover obligation is discharged only if the
    // decode is actually reachable, so an over-strong assumption cannot
    // masquerade as a proof.
    kani::cover!(true, "the decode is reachable under these assumptions");
}

/// **Without the state invariant**: the same decode over a state seeded from
/// arbitrary bytes. `RansDecoder::new` accepts any 4-byte-aligned input, so a
/// malformed or adversarial stream can seed a state at or above 2^63, and
/// `freq * (state >> PRECISION)` is then not guaranteed to fit in u64.
///
/// Kept as a proof harness rather than deleted: if it fails, the failure is the
/// finding — the decoder's totality depends on a precondition the constructor
/// does not enforce, which is exactly the kind of gap a four-target test matrix
/// on well-formed streams cannot see.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(8)]
fn kani_decode_categorical_total_on_arbitrary_stream() {
    let cdf = any_wellformed_cdf3();
    let bytes: [u8; 16] = kani::any();
    let mut d = match RansDecoder::new(&bytes) {
        Ok(d) => d,
        Err(_) => return,
    };
    let s = d.decode_categorical(&cdf);
    kani::assert(
        (s as usize) < cdf.len() - 1,
        "decoded symbol lies inside the alphabet",
    );
    // Non-vacuity, checked rather than asserted. A harness whose assumptions
    // admit no state, or which returns before the call, would satisfy every
    // assertion trivially. This cover obligation is discharged only if the
    // decode is actually reachable, so an over-strong assumption cannot
    // masquerade as a proof.
    kani::cover!(true, "the decode is reachable under these assumptions");
}

/// The **bypass escape** in [`RansDecoder::decode`], over an arbitrary stream.
///
/// When a symbol lands on `max_value` the decoder leaves the CDF and reads raw
/// 4-bit groups: a run-length `n_bypass` accumulated from the stream, then that
/// many nibbles shifted into place by `v << (j * BYPASS_PRECISION)`. Both the
/// run length and the shift distance are therefore *stream-controlled*, and the
/// loop that gathers them reads words without re-checking the buffer.
///
/// This harness asks the verifier whether that path is total for any stream. It
/// carries no assertion of its own: Kani checks arithmetic overflow, shift
/// distance and slice bounds automatically, so whatever it reports is a
/// statement about the decoder rather than about the harness.
#[cfg(kani)]
#[kani::proof]
#[kani::unwind(20)]
fn kani_decode_bypass_escape_is_total() {
    let a: i32 = kani::any();
    kani::assume(0 < a && a < (1i32 << PRECISION));
    let cdf: [i32; 3] = [0, a, 1i32 << PRECISION];
    let bytes: [u8; 64] = kani::any();
    let mut d = match RansDecoder::new(&bytes) {
        Ok(d) => d,
        Err(_) => return,
    };
    let _ = d.decode(&cdf, cdf.len(), 0);
    kani::cover!(true, "decode() returns under these assumptions");
}

/// `RansDecoder::new` is total: every byte slice either yields a decoder or a
/// typed error, and never panics. The constructor is the sink's first contact
/// with untrusted bytes, so its totality is load-bearing on its own.
#[cfg(kani)]
#[kani::proof]
fn kani_decoder_new_is_total() {
    let bytes: [u8; 12] = kani::any();
    let n: usize = kani::any();
    kani::assume(n <= bytes.len());
    match RansDecoder::new(&bytes[..n]) {
        Ok(d) => kani::assert(d.idx == 2, "a seeded decoder starts at word 2"),
        Err(_) => {}
    }
}

/// rANS **encoder** matching [`RansDecoder`] byte-for-byte (rans64, precision 16,
/// bypass 4): the encoder half of the integer codecs' entropy layer. Symbols are
/// buffered in forward order and emitted in reverse on [`flush`](Self::flush), the
/// LIFO rANS discipline. Encode and decode are pure integer, so the round trip is
/// lossless and deterministic on every machine.
///
/// A value whose symbol lies outside the table's span (`value − offset` below 0 or
/// at least `cdf_len − 2`) is coded through the **bypass escape**, exactly as the
/// decoder reads it (`extensions/integer-payloads.md` §2): the escape symbol
/// `cdf_len − 2`; then the count of 4-bit groups, as a run of 15s closed by a group
/// below 15; then the groups of the raw value, least significant first. The raw value
/// is `2·(v − max)` for `v ≥ max` and `2·(−v − 1) + 1` for `v < 0`, where
/// `v = value − offset` and `max = cdf_len − 2`. Each group is a symbol of uniform
/// frequency at precision 4. The encoder's state starts at `2^31` and it writes the
/// words it flushes and no others, so every stream it writes ends canonically
/// ([`RansDecoder::is_canonical_end`]).
#[derive(Default)]
pub struct RansEncoder {
    syms: Vec<Sym>, // forward order
}

/// One coded symbol: a table symbol `(start, freq)` at precision 16, or a 4-bit bypass
/// group.
#[derive(Clone, Copy)]
enum Sym {
    Table { start: u32, freq: u32 },
    Bits(u32),
}

/// The most symbols one value codes as: the table symbol, one count group (a raw value
/// of at most 32 bits has at most eight groups, fewer than 15) and eight value groups.
const MAX_VALUE_SYMS: usize = 10;

/// The symbols `value` codes as against `cdf[0..cdf_len]` with `offset`, in the order
/// the decoder reads them, or `None` where no stream can code it (see
/// [`RansEncoder::encode`]).
fn value_syms(value: i32, cdf: &[i32], cdf_len: usize, offset: i32) -> Option<([Sym; MAX_VALUE_SYMS], usize)> {
    if cdf_len < 3 || cdf_len > cdf.len() {
        return None;
    }
    let max_value = cdf_len as i64 - 2;
    let v = i32::try_from(value as i64 - offset as i64).ok()? as i64;
    let (sym, raw) = if v < 0 {
        (max_value, Some((-2 * v - 1) as u64))
    } else if v >= max_value {
        (max_value, Some(2 * (v - max_value) as u64))
    } else {
        (v, None)
    };
    let (lo, hi) = (cdf[sym as usize], cdf[sym as usize + 1]);
    if lo < 0 || hi <= lo || hi > (1 << PRECISION) {
        return None;
    }
    let mut out = [Sym::Bits(0); MAX_VALUE_SYMS];
    let mut n = 0;
    let mut push = |s: Sym| {
        out[n] = s;
        n += 1;
    };
    push(Sym::Table { start: lo as u32, freq: (hi - lo) as u32 });
    if let Some(raw) = raw {
        // The raw value is at most 2^32 − 1, so at most eight groups.
        let mut n_bypass = 0u32;
        while (raw >> (n_bypass * BYPASS_PRECISION)) != 0 {
            n_bypass += 1;
        }
        let mut count = n_bypass;
        while count >= MAX_BYPASS_VAL {
            push(Sym::Bits(MAX_BYPASS_VAL));
            count -= MAX_BYPASS_VAL;
        }
        push(Sym::Bits(count));
        for j in 0..n_bypass {
            push(Sym::Bits(((raw >> (j * BYPASS_PRECISION)) & MAX_BYPASS_VAL as u64) as u32));
        }
    }
    Some((out, n))
}

/// One step of the encoder, symbols taken last-decoded first: the state takes the
/// symbol, after emitting the low 32 bits while it is too large to.
fn put_sym(state: &mut u64, renorm: &mut Vec<u32>, s: Sym) {
    match s {
        Sym::Table { start, freq } => {
            // renorm so the division stays in range: emit low 32 bits while too big.
            let x_max = ((RANS64_L >> PRECISION) << 32) * freq as u64;
            while *state >= x_max {
                renorm.push(*state as u32);
                *state >>= 32;
            }
            *state = ((*state / freq as u64) << PRECISION) + (*state % freq as u64) + start as u64;
        }
        Sym::Bits(val) => {
            // A group of BYPASS_PRECISION bits is a symbol of frequency
            // 2^(PRECISION − BYPASS_PRECISION): the state takes the bits
            // directly, the inverse of `RansDecoder::get_bits`.
            let freq = 1u64 << (PRECISION - BYPASS_PRECISION);
            let x_max = ((RANS64_L >> PRECISION) << 32) * freq;
            if *state >= x_max {
                renorm.push(*state as u32);
                *state >>= 32;
            }
            *state = (*state << BYPASS_PRECISION) | val as u64;
        }
    }
}

/// The stream: the final state (low word, then high: the decoder seeds `hi << 32 | lo`),
/// then the emitted words in the order the decoder consumes them. Built at its exact
/// length, the words dropped as it is.
fn finish_stream(state: u64, renorm: Vec<u32>) -> Vec<u8> {
    let mut out = Vec::with_capacity(8 + 4 * renorm.len());
    out.extend_from_slice(&(state as u32).to_le_bytes());
    out.extend_from_slice(&((state >> 32) as u32).to_le_bytes());
    for w in renorm.iter().rev() {
        out.extend_from_slice(&w.to_le_bytes());
    }
    out
}

impl RansEncoder {
    pub fn new() -> Self {
        Self::default()
    }

    /// Buffer one symbol: `value` against `cdf[0..cdf_len]` with `offset`, exactly
    /// the arguments [`RansDecoder::decode`] consumes, so the decoder returns `value`.
    /// A symbol outside the table's span is coded through the bypass escape. Returns
    /// `false`, buffering nothing, only where no stream can code `value`: a table
    /// shorter than 3 entries or longer than `cdf`, a frequency that is not positive,
    /// or `value − offset` outside `i32` (the decoder forms the escaped value in
    /// `i32`).
    pub fn encode(&mut self, value: i32, cdf: &[i32], cdf_len: usize, offset: i32) -> bool {
        match value_syms(value, cdf, cdf_len, offset) {
            Some((syms, n)) => {
                self.syms.extend_from_slice(&syms[..n]);
                true
            }
            None => false,
        }
    }

    /// Encode the buffered symbols (in reverse) and return the bitstream that
    /// [`RansDecoder`] decodes back to the same symbols in forward order.
    pub fn flush(self) -> Vec<u8> {
        let mut state: u64 = RANS64_L;
        let mut renorm: Vec<u32> = Vec::new();
        for &s in self.syms.iter().rev() {
            put_sym(&mut state, &mut renorm, s);
        }
        finish_stream(state, renorm)
    }
}

/// The same encoder, fed values in reverse: each [`put`](Self::put) codes a value the
/// decoder reads *before* every value put so far. It holds the state and the words it
/// has emitted, never a buffer of symbols, so a caller that can produce its values
/// last-first (a two-pass encoder) codes a stream in about the stream's own size. Fed a
/// sequence of values last-first, it writes exactly the bytes [`RansEncoder`] writes for
/// the sequence.
pub struct RansReverseEncoder {
    state: u64,
    renorm: Vec<u32>,
}

impl Default for RansReverseEncoder {
    fn default() -> Self {
        RansReverseEncoder { state: RANS64_L, renorm: Vec::new() }
    }
}

impl RansReverseEncoder {
    pub fn new() -> Self {
        Self::default()
    }

    /// Code `value` ahead of every value put so far, against `cdf[0..cdf_len]` with
    /// `offset`; `false`, coding nothing, exactly where [`RansEncoder::encode`] is
    /// `false`.
    pub fn put(&mut self, value: i32, cdf: &[i32], cdf_len: usize, offset: i32) -> bool {
        match value_syms(value, cdf, cdf_len, offset) {
            Some((syms, n)) => {
                for &s in syms[..n].iter().rev() {
                    put_sym(&mut self.state, &mut self.renorm, s);
                }
                true
            }
            None => false,
        }
    }

    /// The words emitted so far.
    pub fn words(&self) -> usize {
        self.renorm.len()
    }

    /// The bitstream.
    pub fn finish(self) -> Vec<u8> {
        finish_stream(self.state, self.renorm)
    }
}

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

    /// Quantized zero-mean Gaussian CDF for scale `sigma` over `n` real symbols
    /// (centered), plus the decoder's escape slot. `cdf_len = n + 2`; freqs sum
    /// to 2^PRECISION, each real symbol ≥ 1 so it is codable.
    fn gauss_cdf(sigma: f64, n: usize) -> Vec<i32> {
        let total = 1i32 << PRECISION;
        let center = n as f64 / 2.0;
        let pmf: Vec<f64> = (0..n).map(|i| {
            let d = i as f64 - center;
            (-d * d / (2.0 * sigma * sigma)).exp()
        }).collect();
        let s: f64 = pmf.iter().sum();
        // distribute total-1 over the n real symbols (≥1 each), 1 to the escape.
        let mut f: Vec<i32> = pmf.iter().map(|&p| (((p / s) * (total - 1) as f64).round() as i32).max(1)).collect();
        let sum: i32 = f.iter().sum();
        let amax = (0..n).max_by_key(|&i| f[i]).unwrap();
        f[amax] += (total - 1) - sum; // real symbols now sum to total-1
        f.push(1); // escape freq
        let mut cdf = vec![0i32; n + 2];
        for i in 0..n + 1 { cdf[i + 1] = cdf[i] + f[i]; }
        cdf // cdf[n+1] = total
    }

    #[test]
    fn rans_encode_decode_round_trips() {
        let cdf = gauss_cdf(3.0, 16);
        let cdf_len = cdf.len();
        let offset = -8; // symbols centered around 0
        // a deterministic pseudo-random symbol stream in range
        let syms: Vec<i32> = (0..500).map(|i| ((i * 2654435761usize) % 16) as i32 + offset).collect();
        let mut enc = RansEncoder::new();
        for &v in &syms {
            assert!(enc.encode(v, &cdf, cdf_len, offset), "symbol {v} out of range");
        }
        let bits = enc.flush();
        let mut dec = RansDecoder::new(&bits).unwrap();
        let out: Vec<i32> = (0..syms.len()).map(|_| dec.decode(&cdf, cdf_len, offset)).collect();
        assert_eq!(out, syms, "rANS round-trip must be lossless");
    }

    /// The reverse writer, fed the values last-first, writes the bytes the buffering
    /// encoder writes, table symbols and bypass escapes alike, and refuses what it
    /// refuses.
    #[test]
    fn the_reverse_writer_writes_the_same_bytes() {
        let mut x: u64 = 0x2545_f491_4f6c_dd1d;
        let mut next = || {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            x
        };
        for (n, offset, count) in [(1usize, 0i32, 0usize), (1, 0, 1), (4, -2, 300), (16, -8, 5000), (255, 7, 3000), (511, -255, 20000)] {
            let cdf = gauss_cdf(n as f64 / 6.0 + 0.5, n);
            let cdf_len = cdf.len();
            let vals: Vec<i32> = (0..count)
                .map(|_| {
                    let r = next();
                    match r % 8 {
                        0 => r as i32,
                        1 => ((r >> 8) as i32) >> ((r >> 40) % 31),
                        _ => offset + ((r >> 8) % (n as u64 + 2)) as i32 - 1,
                    }
                })
                .collect();
            let mut fwd = RansEncoder::new();
            let fits: Vec<bool> = vals.iter().map(|&v| fwd.encode(v, &cdf, cdf_len, offset)).collect();
            let mut rev = RansReverseEncoder::new();
            for (i, &v) in vals.iter().enumerate().rev() {
                assert_eq!(rev.put(v, &cdf, cdf_len, offset), fits[i], "value {v}");
            }
            assert_eq!(rev.finish(), fwd.flush(), "n={n} offset={offset} count={count}");
        }
    }

    /// The bypass escape round-trips every `i32` the decoder can return, against
    /// tables of several lengths and offsets, the boundaries `±2^31` included, and
    /// every stream ends canonically. A value no stream can code (`value − offset`
    /// outside `i32`) is refused and buffers nothing.
    #[test]
    fn encoder_bypass_roundtrip() {
        let mut x: u64 = 0x9e37_79b9_7f4a_7c15;
        let mut next = || {
            x ^= x << 13;
            x ^= x >> 7;
            x ^= x << 17;
            x
        };
        for (n, offset) in [(1usize, 0i32), (4, -2), (16, -8), (255, 7)] {
            let cdf = gauss_cdf(2.0, n);
            let cdf_len = cdf.len();
            let mut vals: Vec<i32> = vec![i32::MIN, i32::MAX, i32::MIN + 1, i32::MAX - 1, -1, 0, 1, offset, offset - 1,
                                          offset + n as i32, offset + n as i32 - 1, offset + n as i32 + 1];
            for _ in 0..400 {
                let r = next();
                vals.push(match r % 4 {
                    0 => r as i32,
                    1 => ((r >> 8) % (n as u64 + 4)) as i32 + offset - 2,
                    2 => ((r >> 8) as i32) >> ((r >> 40) % 31),
                    _ => offset + ((r >> 8) % n as u64) as i32,
                });
            }
            let mut enc = RansEncoder::new();
            let mut coded = Vec::new();
            for &v in &vals {
                let fits = i32::try_from(v as i64 - offset as i64).is_ok();
                assert_eq!(enc.encode(v, &cdf, cdf_len, offset), fits, "value {v} offset {offset}");
                if fits {
                    coded.push(v);
                }
            }
            let bits = enc.flush();
            let mut dec = RansDecoder::new(&bits).unwrap();
            let out: Vec<i32> = (0..coded.len()).map(|_| dec.decode(&cdf, cdf_len, offset)).collect();
            assert_eq!(out, coded, "bypass round-trip (n={n}, offset={offset})");
            assert!(dec.is_canonical_end(), "the encoder's stream ends canonically (n={n})");
        }
        // A table too short to have an escape, or a length past the slice, codes nothing.
        let mut enc = RansEncoder::new();
        assert!(!enc.encode(0, &[0, 65536], 2, 0));
        assert!(!enc.encode(0, &[0, 1, 65536], 4, 0));
        assert_eq!(enc.flush(), RansEncoder::new().flush());
    }

    #[test]
    fn conditional_integer_entropy_round_trips_via_int_entropy() {
        // The codec's entropy LAYER end-to-end, pure integer: each latent's CDF
        // is chosen by its scale σ via the integer `int_entropy` bucketing
        // (det_build_index), encoded, then decoded back identically. This is what
        // makes a learned conditional codec cross-hardware exact.
        use crate::int_entropy::{det_build_index, fixed_scale_edges, quantize_scale};
        const FRAC: u32 = 12;
        let scale_table: Vec<f32> = vec![1.0, 2.0, 3.5, 6.0, 10.0]; // σ buckets
        let edges = fixed_scale_edges(&scale_table, FRAC);
        let n = 24usize;
        let offset = -(n as i32 / 2);
        let cdfs: Vec<Vec<i32>> = scale_table.iter().map(|&s| gauss_cdf(s as f64, n)).collect();

        // (latent value, its scale σ) — deterministic stream, kept in range.
        let stream: Vec<(i32, f32)> = (0..400).map(|i| {
            let sigma = scale_table[(i * 7) % scale_table.len()];
            let v = ((i * 40503) % n) as i32 + offset;
            (v, sigma)
        }).collect();

        let mut enc = RansEncoder::new();
        let mut buckets = Vec::new();
        for &(v, sigma) in &stream {
            let b = det_build_index(quantize_scale(sigma, FRAC), &edges); // INTEGER σ→bucket
            buckets.push(b);
            let cdf = &cdfs[b];
            assert!(enc.encode(v, cdf, cdf.len(), offset), "latent {v} out of range for σ={sigma}");
        }
        let bits = enc.flush();

        // Decode: recompute the SAME bucket from σ (the encoder and decoder agree
        // because the bucketing is integer-deterministic) and decode each latent.
        let mut dec = RansDecoder::new(&bits).unwrap();
        for (i, &(v, sigma)) in stream.iter().enumerate() {
            let b = det_build_index(quantize_scale(sigma, FRAC), &edges);
            assert_eq!(b, buckets[i], "bucket must match encoder");
            let cdf = &cdfs[b];
            assert_eq!(dec.decode(cdf, cdf.len(), offset), v, "conditional decode mismatch at {i}");
        }
    }
}
