669 lines
29 KiB
Rust
669 lines
29 KiB
Rust
//! 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}");
|
||
}
|
||
}
|
||
}
|