Skip to main content

ryg_rans_rs_core/
lib.rs

1#![no_std]
2#![forbid(unsafe_code)]
3
4//! # ryg-rans-rs-core
5//!
6//! **Deterministic algorithmic core of rANS entropy coding.**
7//! `#![no_std]` · `#![forbid(unsafe_code)]` · zero-allocation encode/decode hot paths.
8//!
9//! This crate provides a complete native Rust reconstruction of both the 32-bit
10//! byte-aligned and 64-bit rANS encoder/decoder variants from Fabian Giesen's
11//! public-domain [`ryg_rans`](https://github.com/rygorous/ryg_rans) repository.
12//!
13//! ## Variants
14//!
15//! | Variant | Upstream File | Key Constants | Renormalization Unit | State Width |
16//! |---------|---------------|---------------|---------------------|-------------|
17//! | Byte rANS (32-bit) | `rans_byte.h` | `RANS_BYTE_L = 2^23` | 8-bit (byte) | 31-bit effective |
18//! | 64-bit rANS | `rans64.h` | `RANS64_L = 2^31` | 32-bit (word) | 63-bit effective |
19//!
20//! Both variants provide:
21//! - **Division-based reference path**: `C(s,x) = ((x/freq) << scale_bits) + (x%freq) + start`
22//! - **Reciprocal-multiply fast path**: Uses multiply-high approximation to avoid
23//!   integer division in the encode hot loop (Alverson's method).
24//! - **Two-state interleaving**: Encodes symbols into two interleaved streams
25//!   for superscalar throughput.
26//! - **Step-only operations**: Decoder advance without renormalization, enabling
27//!   interleaved decoding patterns.
28//!
29//! ## Encoding Semantics (from upstream)
30//!
31//! - **Reverse order**: Symbols must be encoded last-to-first (stack discipline).
32//! - **Reverse-growing output**: The backward writer starts at the end of the
33//!   buffer and moves toward the beginning.
34//! - **Renormalization**: When the state exceeds a symbol-specific threshold
35//!   `x_max`, the lowest byte/word is emitted and the state is shifted.
36//! - **Flush**: The remaining state is written directly as 4 bytes (byte rANS)
37//!   or 2 × u32 words (64-bit rANS).
38//!
39//! ## Error Handling
40//!
41//! - **Encoding**: All encode operations return `Result<(), EncodeError>`. The
42//!   only error variant is `OutputTooSmall`, which occurs when the backward
43//!   writer's buffer is exhausted. State mutation is transactional — if the
44//!   error occurs, the state is left in a consistent but partially-advanced
45//!   position.
46//! - **Decoding**: All decode operations return `Result<(), DecodeError>`.
47//!   `InputTooShort` indicates a truncated compressed stream.
48//!
49//! ## I/O Abstraction
50//!
51//! Trait-based readers and writers allow the core algorithms to operate on any
52//! storage backend. Concrete implementations provided:
53//! - `BackwardByteWriter` / `BackwardWord32Writer` for encoding output
54//! - `ByteReader` / `Word32Reader` for decoding input
55//! - `SliceBackwardWriter` for convenient `&mut [u8]` encoding
56//! - `&[u8]` implements `ForwardReader` directly for decoding
57
58use core::fmt;
59
60pub mod malformed;
61
62// ---------------------------------------------------------------------------
63// Error types
64// ---------------------------------------------------------------------------
65
66/// Errors that can occur during encoding.
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68pub enum EncodeError {
69    /// The output buffer is too small to hold the encoded data.
70    OutputTooSmall,
71}
72
73impl fmt::Display for EncodeError {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        match self {
76            EncodeError::OutputTooSmall => write!(f, "output buffer too small"),
77        }
78    }
79}
80
81/// Errors that can occur during decoding.
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub enum DecodeError {
84    /// The input is too short for the requested operation (truncated stream).
85    InputTooShort,
86}
87
88impl fmt::Display for DecodeError {
89    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
90        match self {
91            DecodeError::InputTooShort => write!(f, "truncated input stream"),
92        }
93    }
94}
95
96// ---------------------------------------------------------------------------
97// Symbol construction error types
98// ---------------------------------------------------------------------------
99
100/// Errors that can occur during symbol construction.
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub enum ModelError {
103    /// The input sequence was empty.
104    EmptyInput,
105    /// Total frequency is zero.
106    ZeroTotal,
107    /// `scale_bits` is outside the valid range.
108    InvalidScaleBits,
109    /// Symbol frequency is zero.
110    ZeroFrequency,
111    /// Frequency exceeds the allowed range for the given scale_bits/start.
112    FrequencyOutOfRange,
113    /// Start value exceeds the allowed range.
114    StartOutOfRange,
115    /// The provided total does not match the accumulated total.
116    TotalMismatch,
117    /// The workspace buffer is too small for the requested operation.
118    WorkspaceTooSmall,
119}
120
121impl fmt::Display for ModelError {
122    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
123        match self {
124            ModelError::EmptyInput => write!(f, "empty input sequence"),
125            ModelError::ZeroTotal => write!(f, "total frequency is zero"),
126            ModelError::InvalidScaleBits => write!(f, "scale_bits is out of valid range"),
127            ModelError::ZeroFrequency => write!(f, "symbol frequency is zero"),
128            ModelError::FrequencyOutOfRange => write!(f, "frequency exceeds allowed range"),
129            ModelError::StartOutOfRange => write!(f, "start value exceeds allowed range"),
130            ModelError::TotalMismatch => write!(f, "total frequency mismatch"),
131            ModelError::WorkspaceTooSmall => write!(f, "workspace buffer too small"),
132        }
133    }
134}
135
136#[cfg(feature = "std")]
137extern crate std;
138
139#[cfg(feature = "alloc")]
140extern crate alloc;
141
142#[cfg(feature = "std")]
143impl std::error::Error for EncodeError {}
144
145#[cfg(feature = "std")]
146impl std::error::Error for DecodeError {}
147
148#[cfg(feature = "std")]
149impl std::error::Error for ModelError {}
150
151/// Lower bound of the normalization interval.
152///
153/// Equivalent to `RANS_BYTE_L` in `rans_byte.h`.
154pub const RANS_BYTE_L: u32 = 1u32 << 23;
155
156/// 32-bit rANS encoder/decoder state.
157///
158/// Equivalent to `RansState` (typedef for `uint32_t`).
159#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
160pub struct RansByteState(pub u32);
161
162impl RansByteState {
163    /// Create a new state initialized to the lower bound.
164    #[inline]
165    pub const fn new() -> Self {
166        Self(RANS_BYTE_L)
167    }
168
169    /// Return the raw state value.
170    #[inline]
171    pub const fn get(&self) -> u32 {
172        self.0
173    }
174}
175
176impl Default for RansByteState {
177    #[inline]
178    fn default() -> Self {
179        Self::new()
180    }
181}
182
183// ---------------------------------------------------------------------------
184// Backward byte writer
185// ---------------------------------------------------------------------------
186
187/// A backward-growing byte writer.
188///
189/// Starts at the logical end of the provided buffer and writes bytes toward
190/// the beginning. Used for rANS encoding output.
191///
192/// Equivalent to `uint8_t* ptr` pointing to the end of the output buffer
193/// in the upstream C code.
194pub struct BackwardByteWriter<'a> {
195    buf: &'a mut [u8],
196    pos: usize, // current write position (0 <= pos <= len)
197}
198
199impl<'a> BackwardByteWriter<'a> {
200    /// Create a new backward writer from a mutable byte slice.
201    ///
202    /// The writer starts at the end of the buffer. `pos` tracks the
203    /// zero-based index; initially `pos == len`.
204    #[inline]
205    pub fn new(buf: &'a mut [u8]) -> Self {
206        let len = buf.len();
207        Self { buf, pos: len }
208    }
209
210    /// Write a single byte at the current position (decrementing the cursor).
211    ///
212    /// Returns `Ok(())` if there was room, `Err(())` if the writer has
213    /// exhausted its buffer.
214    #[inline]
215    pub fn write_byte(&mut self, b: u8) -> Result<(), ()> {
216        if self.pos == 0 {
217            return Err(());
218        }
219        self.pos -= 1;
220        self.buf[self.pos] = b;
221        Ok(())
222    }
223
224    /// Write 4 bytes (little-endian u32) at the current position.
225    ///
226    /// Returns `Ok(())` if there was room, `Err(())` otherwise.
227    #[inline]
228    pub fn write_u32_le(&mut self, v: u32) -> Result<(), ()> {
229        if self.pos < 4 {
230            return Err(());
231        }
232        self.pos -= 4;
233        self.buf[self.pos..self.pos + 4].copy_from_slice(&v.to_le_bytes());
234        Ok(())
235    }
236
237    /// Current zero-based write position (number of bytes before the cursor).
238    #[inline]
239    pub fn position(&self) -> usize {
240        self.pos
241    }
242
243    /// Number of bytes written so far.
244    #[inline]
245    pub fn bytes_written(&self) -> usize {
246        self.buf.len() - self.pos
247    }
248
249    /// Return the encoded portion (from current position to end) as a slice.
250    #[inline]
251    pub fn encoded(&self) -> &[u8] {
252        &self.buf[self.pos..]
253    }
254
255    /// Return the remaining capacity (number of bytes that can still be written).
256    #[inline]
257    pub fn remaining(&self) -> usize {
258        self.pos
259    }
260}
261
262impl fmt::Debug for BackwardByteWriter<'_> {
263    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
264        f.debug_struct("BackwardByteWriter")
265            .field("pos", &self.pos)
266            .field("len", &self.buf.len())
267            .finish()
268    }
269}
270
271// ---------------------------------------------------------------------------
272// Forward byte reader
273// ---------------------------------------------------------------------------
274
275/// A forward-growing byte reader.
276///
277/// Used for rANS decoding input. Reads bytes from the beginning of the
278/// compressed stream moving forward.
279///
280/// Equivalent to `uint8_t* ptr` pointing to the start of compressed data
281/// in the upstream C code.
282#[derive(Clone)]
283pub struct ByteReader<'a> {
284    buf: &'a [u8],
285    pos: usize,
286}
287
288impl<'a> ByteReader<'a> {
289    /// Create a new byte reader from a byte slice.
290    #[inline]
291    pub fn new(buf: &'a [u8]) -> Self {
292        Self { buf, pos: 0 }
293    }
294
295    /// Read a single byte at the current position, advancing forward.
296    ///
297    /// Returns `None` if the buffer is exhausted.
298    #[inline]
299    pub fn read_byte(&mut self) -> Option<u8> {
300        if self.pos >= self.buf.len() {
301            return None;
302        }
303        let b = self.buf[self.pos];
304        self.pos += 1;
305        Some(b)
306    }
307
308    /// Read 4 bytes as a little-endian u32.
309    ///
310    /// Returns `None` if fewer than 4 bytes remain.
311    #[inline]
312    pub fn read_u32_le(&mut self) -> Option<u32> {
313        if self.pos + 4 > self.buf.len() {
314            return None;
315        }
316        let v = u32::from_le_bytes([
317            self.buf[self.pos],
318            self.buf[self.pos + 1],
319            self.buf[self.pos + 2],
320            self.buf[self.pos + 3],
321        ]);
322        self.pos += 4;
323        Some(v)
324    }
325
326    /// Current read position.
327    #[inline]
328    pub fn position(&self) -> usize {
329        self.pos
330    }
331
332    /// Number of bytes consumed so far.
333    #[inline]
334    pub fn bytes_consumed(&self) -> usize {
335        self.pos
336    }
337
338    /// Remaining unread bytes.
339    #[inline]
340    pub fn remaining(&self) -> usize {
341        self.buf.len().saturating_sub(self.pos)
342    }
343}
344
345impl fmt::Debug for ByteReader<'_> {
346    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
347        f.debug_struct("ByteReader")
348            .field("pos", &self.pos)
349            .field("len", &self.buf.len())
350            .finish()
351    }
352}
353
354// ---------------------------------------------------------------------------
355// Encoder symbol (reciprocal-multiply setup)
356// ---------------------------------------------------------------------------
357
358/// Precomputed reciprocal encoder symbol.
359///
360/// Equivalent to `RansEncSymbol` in `rans_byte.h`.
361///
362/// The fast encoder uses a multiply-high approximation to divide by the
363/// symbol frequency, avoiding an integer division in the hot loop.
364#[derive(Debug, Clone, Copy)]
365pub struct RansByteEncSymbol {
366    /// (Exclusive) upper bound of pre-normalization interval.
367    pub x_max: u32,
368    /// Fixed-point reciprocal frequency.
369    pub rcp_freq: u32,
370    /// Bias value.
371    pub bias: u32,
372    /// Complement of frequency: `(1 << scale_bits) - freq`.
373    pub cmpl_freq: u16,
374    /// Reciprocal shift amount.
375    pub rcp_shift: u16,
376}
377
378impl RansByteEncSymbol {
379    /// Initialize an encoder symbol with validation.
380    ///
381    /// Equivalent to `RansEncSymbolInit` in `rans_byte.h`.
382    ///
383    /// Returns `Err(ModelError::InvalidScaleBits)` if `scale_bits` is not in
384    /// `1..=16`. Returns `Err(ModelError::ZeroFrequency)` if `freq == 0`.
385    /// Returns `Err(ModelError::StartOutOfRange)` if `start > (1 << scale_bits)`.
386    /// Returns `Err(ModelError::FrequencyOutOfRange)` if
387    /// `freq > (1 << scale_bits) - start`.
388    #[inline]
389    pub fn new(start: u32, freq: u32, scale_bits: u32) -> Result<Self, ModelError> {
390        if !(1..=16).contains(&scale_bits) {
391            return Err(ModelError::InvalidScaleBits);
392        }
393        let max_start = 1u64 << scale_bits;
394        if (start as u64) > max_start {
395            return Err(ModelError::StartOutOfRange);
396        }
397        if freq == 0 {
398            return Err(ModelError::ZeroFrequency);
399        }
400        if (freq as u64) > max_start - (start as u64) {
401            return Err(ModelError::FrequencyOutOfRange);
402        }
403
404        Ok(Self::new_unchecked(start, freq, scale_bits))
405    }
406
407    /// Initialize an encoder symbol without validation.
408    ///
409    /// Caller must ensure preconditions are met. Prefer [`new`] for
410    /// external use.
411    #[inline]
412    pub(crate) fn new_unchecked(start: u32, freq: u32, scale_bits: u32) -> Self {
413        debug_assert!(scale_bits <= 16, "scale_bits must be <= 16");
414        debug_assert!(start <= (1u32 << scale_bits), "start out of range");
415        debug_assert!(freq <= (1u32 << scale_bits) - start, "freq out of range");
416
417        let x_max = ((RANS_BYTE_L >> scale_bits) << 8) * freq;
418        let cmpl_freq = ((1u32 << scale_bits) - freq) as u16;
419
420        if freq < 2 {
421            // freq == 1 special case
422            // rcp_freq = ~0u, rcp_shift = 0
423            // bias = start + (1 << scale_bits) - 1
424            Self {
425                x_max,
426                rcp_freq: !0u32,
427                rcp_shift: 0,
428                bias: start + (1u32 << scale_bits) - 1,
429                cmpl_freq,
430            }
431        } else {
432            // Alverson "Integer Division using reciprocals"
433            let mut shift = 0u32;
434            while freq > (1u32 << shift) {
435                shift += 1;
436            }
437
438            // rcp_freq = (((1ull << (shift + 31)) + freq - 1) / freq) as u32
439            let rcp_freq = (((1u64 << (shift + 31)) + freq as u64 - 1) / freq as u64) as u32;
440            let rcp_shift = shift - 1;
441
442            Self {
443                x_max,
444                rcp_freq,
445                rcp_shift: rcp_shift as u16,
446                bias: start,
447                cmpl_freq,
448            }
449        }
450    }
451}
452
453// ---------------------------------------------------------------------------
454// Decoder symbol
455// ---------------------------------------------------------------------------
456
457/// Decoder symbol description.
458///
459/// Equivalent to `RansDecSymbol` in `rans_byte.h`.
460#[derive(Debug, Clone, Copy, PartialEq, Eq)]
461pub struct RansByteDecSymbol {
462    /// Start of range.
463    pub start: u16,
464    /// Symbol frequency.
465    pub freq: u16,
466}
467
468impl RansByteDecSymbol {
469    /// Initialize a decoder symbol with validation.
470    ///
471    /// Equivalent to `RansDecSymbolInit`.
472    ///
473    /// Returns `Err(ModelError::ZeroFrequency)` if `freq == 0`.
474    /// Returns `Err(ModelError::StartOutOfRange)` if `start > (1 << 16)`.
475    /// Returns `Err(ModelError::FrequencyOutOfRange)` if
476    /// `freq > (1 << 16) - start`.
477    #[inline]
478    pub fn new(start: u32, freq: u32) -> Result<Self, ModelError> {
479        if freq == 0 {
480            return Err(ModelError::ZeroFrequency);
481        }
482        if (start as u64) > (1u64 << 16) {
483            return Err(ModelError::StartOutOfRange);
484        }
485        if (freq as u64) > (1u64 << 16) - (start as u64) {
486            return Err(ModelError::FrequencyOutOfRange);
487        }
488        Ok(Self::new_unchecked(start, freq))
489    }
490
491    /// Initialize a decoder symbol without validation.
492    ///
493    /// Caller must ensure preconditions are met. Prefer [`new`] for
494    /// external use.
495    #[inline]
496    pub(crate) fn new_unchecked(start: u32, freq: u32) -> Self {
497        debug_assert!(start <= (1u32 << 16), "start out of range");
498        debug_assert!(freq <= (1u32 << 16) - start, "freq out of range");
499        Self {
500            start: start as u16,
501            freq: freq as u16,
502        }
503    }
504}
505
506// ---------------------------------------------------------------------------
507// Division-based encoder (reference path)
508// ---------------------------------------------------------------------------
509
510/// Renormalize the encoder state.
511///
512/// Equivalent to `RansEncRenorm` in `rans_byte.h`.
513///
514/// Emits bytes (LSB first) until the state falls below the threshold.
515/// Returns `Err(EncodeError::OutputTooSmall)` if the output buffer is exhausted.
516#[inline]
517pub fn rans_byte_enc_renorm<W: BackwardWriter>(
518    x: u32,
519    writer: &mut W,
520    freq: u32,
521    scale_bits: u32,
522) -> Result<u32, EncodeError> {
523    let x_max = ((RANS_BYTE_L >> scale_bits) << 8) * freq;
524    let mut x = x;
525    if x >= x_max {
526        while x >= x_max {
527            writer
528                .write_byte((x & 0xff) as u8)
529                .map_err(|_| EncodeError::OutputTooSmall)?;
530            x >>= 8;
531        }
532    }
533    Ok(x)
534}
535
536/// Encoder put-symbol (division-based reference path).
537///
538/// Equivalent to `RansEncPut` in `rans_byte.h`.
539/// Returns `Err(EncodeError::OutputTooSmall)` if the output buffer is exhausted.
540#[inline]
541pub fn rans_byte_enc_put<W: BackwardWriter>(
542    state: &mut RansByteState,
543    writer: &mut W,
544    start: u32,
545    freq: u32,
546    scale_bits: u32,
547) -> Result<(), EncodeError> {
548    let x = rans_byte_enc_renorm(state.0, writer, freq, scale_bits)?;
549    state.0 = ((x / freq) << scale_bits) + (x % freq) + start;
550    Ok(())
551}
552
553/// Flush the encoder.
554///
555/// Equivalent to `RansEncFlush` in `rans_byte.h`.
556///
557/// Writes the full 32-bit state in little-endian order (4 bytes).
558/// Returns `Err(EncodeError::OutputTooSmall)` if the output buffer is exhausted.
559#[inline]
560pub fn rans_byte_enc_flush<W: BackwardWriter>(
561    state: &RansByteState,
562    writer: &mut W,
563) -> Result<(), EncodeError> {
564    let x = state.0;
565    writer
566        .write_u32_le(x)
567        .map_err(|_| EncodeError::OutputTooSmall)
568}
569
570// ---------------------------------------------------------------------------
571// Division-based decoder (reference path)
572// ---------------------------------------------------------------------------
573
574/// Initialize the decoder.
575///
576/// Equivalent to `RansDecInit` in `rans_byte.h`.
577///
578/// Reads 4 bytes as the initial state.
579#[inline]
580pub fn rans_byte_dec_init<R: ForwardReader>(reader: &mut R) -> Result<RansByteState, DecodeError> {
581    let x = reader.read_u32_le().ok_or(DecodeError::InputTooShort)?;
582    Ok(RansByteState(x))
583}
584
585/// Get the cumulative frequency from the current state.
586///
587/// Equivalent to `RansDecGet` in `rans_byte.h`.
588#[inline]
589pub fn rans_byte_dec_get(state: &RansByteState, scale_bits: u32) -> u32 {
590    state.0 & ((1u32 << scale_bits) - 1)
591}
592
593/// Advance the decoder by a symbol (with renormalization).
594///
595/// Equivalent to `RansDecAdvance` in `rans_byte.h`.
596#[inline]
597pub fn rans_byte_dec_advance<R: ForwardReader>(
598    state: &mut RansByteState,
599    reader: &mut R,
600    start: u32,
601    freq: u32,
602    scale_bits: u32,
603) -> Result<(), DecodeError> {
604    let mask = (1u32 << scale_bits) - 1;
605    let x = state.0;
606    let mut x = freq * (x >> scale_bits) + (x & mask) - start;
607
608    // renormalize
609    if x < RANS_BYTE_L {
610        loop {
611            let b = reader.read_byte().ok_or(DecodeError::InputTooShort)?;
612            x = (x << 8) | (b as u32);
613            if x >= RANS_BYTE_L {
614                break;
615            }
616        }
617    }
618
619    state.0 = x;
620    Ok(())
621}
622
623// ---------------------------------------------------------------------------
624// Reciprocal-multiply encoder (fast path)
625// ---------------------------------------------------------------------------
626
627/// Encoder put-symbol using precomputed reciprocal (fast path).
628///
629/// Equivalent to `RansEncPutSymbol` in `rans_byte.h`.
630///
631/// This is the fast path that avoids integer division in the hot loop.
632/// Returns `Err(EncodeError::OutputTooSmall)` if the output buffer is exhausted.
633#[inline]
634pub fn rans_byte_enc_put_symbol<W: BackwardWriter>(
635    state: &mut RansByteState,
636    writer: &mut W,
637    sym: &RansByteEncSymbol,
638) -> Result<(), EncodeError> {
639    debug_assert!(sym.x_max != 0, "cannot encode symbol with freq=0");
640
641    let mut x = state.0;
642    if x >= sym.x_max {
643        while x >= sym.x_max {
644            writer
645                .write_byte((x & 0xff) as u8)
646                .map_err(|_| EncodeError::OutputTooSmall)?;
647            x >>= 8;
648        }
649    }
650
651    // x = C(s,x)
652    // 32-bit "multiply high": (u64)x * rcp_freq >> 32 >> rcp_shift
653    let q = (((x as u64) * (sym.rcp_freq as u64)) >> 32) >> sym.rcp_shift;
654    state.0 = x + sym.bias + (q as u32) * (sym.cmpl_freq as u32);
655    Ok(())
656}
657
658// ---------------------------------------------------------------------------
659// Decoder advance with symbol (convenience)
660// ---------------------------------------------------------------------------
661
662/// Advance the decoder using a decoder symbol.
663///
664/// Equivalent to `RansDecAdvanceSymbol` in `rans_byte.h`.
665#[inline]
666pub fn rans_byte_dec_advance_symbol<R: ForwardReader>(
667    state: &mut RansByteState,
668    reader: &mut R,
669    sym: &RansByteDecSymbol,
670    scale_bits: u32,
671) -> Result<(), DecodeError> {
672    rans_byte_dec_advance(state, reader, sym.start as u32, sym.freq as u32, scale_bits)
673}
674
675// ---------------------------------------------------------------------------
676// Step-only operations (for interleaving)
677// ---------------------------------------------------------------------------
678
679/// Advance the decoder without renormalization (interleaving step).
680///
681/// Equivalent to `RansDecAdvanceStep` in `rans_byte.h`.
682#[inline]
683pub fn rans_byte_dec_advance_step(
684    state: &mut RansByteState,
685    start: u32,
686    freq: u32,
687    scale_bits: u32,
688) {
689    let mask = (1u32 << scale_bits) - 1;
690    let x = state.0;
691    state.0 = freq * (x >> scale_bits) + (x & mask) - start;
692}
693
694/// Advance the decoder step using a decoder symbol.
695///
696/// Equivalent to `RansDecAdvanceSymbolStep`.
697#[inline]
698pub fn rans_byte_dec_advance_symbol_step(
699    state: &mut RansByteState,
700    sym: &RansByteDecSymbol,
701    scale_bits: u32,
702) {
703    rans_byte_dec_advance_step(state, sym.start as u32, sym.freq as u32, scale_bits);
704}
705
706/// Decoder renormalization only.
707///
708/// Equivalent to `RansDecRenorm` in `rans_byte.h`.
709#[inline]
710pub fn rans_byte_dec_renorm<R: ForwardReader>(
711    state: &mut RansByteState,
712    reader: &mut R,
713) -> Result<(), DecodeError> {
714    let mut x = state.0;
715    if x < RANS_BYTE_L {
716        loop {
717            let b = reader.read_byte().ok_or(DecodeError::InputTooShort)?;
718            x = (x << 8) | (b as u32);
719            if x >= RANS_BYTE_L {
720                break;
721            }
722        }
723        state.0 = x;
724    }
725    Ok(())
726}
727
728// Advance
729// ---------------------------------------------------------------------------
730// BackwardWriter / ForwardReader traits
731// ---------------------------------------------------------------------------
732
733/// Trait for backward-writing output (encoder output).
734pub trait BackwardWriter {
735    /// Write a single byte.
736    fn write_byte(&mut self, b: u8) -> Result<(), ()>;
737
738    /// Write a u32 in little-endian format (4 bytes).
739    fn write_u32_le(&mut self, v: u32) -> Result<(), ()>;
740}
741
742/// Trait for forward-reading input (decoder input).
743pub trait ForwardReader {
744    /// Read a single byte.
745    fn read_byte(&mut self) -> Option<u8>;
746
747    /// Read a u32 in little-endian format (4 bytes).
748    fn read_u32_le(&mut self) -> Option<u32>;
749}
750
751// Implement the traits for our concrete types.
752impl<'a> BackwardWriter for BackwardByteWriter<'a> {
753    #[inline]
754    fn write_byte(&mut self, b: u8) -> Result<(), ()> {
755        self.write_byte(b)
756    }
757
758    #[inline]
759    fn write_u32_le(&mut self, v: u32) -> Result<(), ()> {
760        self.write_u32_le(v)
761    }
762}
763
764impl<'a> ForwardReader for ByteReader<'a> {
765    #[inline]
766    fn read_byte(&mut self) -> Option<u8> {
767        self.read_byte()
768    }
769
770    #[inline]
771    fn read_u32_le(&mut self) -> Option<u32> {
772        self.read_u32_le()
773    }
774}
775
776/// A wrapper around `&mut [u8]` that implements `BackwardWriter`.
777///
778/// This avoids issues with implementing a trait on a mutable reference to an unsized slice.
779pub struct SliceBackwardWriter<'a>(pub &'a mut [u8]);
780
781impl BackwardWriter for SliceBackwardWriter<'_> {
782    #[inline]
783    fn write_byte(&mut self, b: u8) -> Result<(), ()> {
784        let buf = core::mem::take(&mut self.0);
785        if buf.is_empty() {
786            self.0 = buf;
787            return Err(());
788        }
789        let len = buf.len();
790        buf[len - 1] = b;
791        self.0 = &mut buf[..len - 1];
792        Ok(())
793    }
794
795    #[inline]
796    fn write_u32_le(&mut self, v: u32) -> Result<(), ()> {
797        let buf = core::mem::take(&mut self.0);
798        let len = buf.len();
799        if len < 4 {
800            self.0 = buf;
801            return Err(());
802        }
803        let bytes = v.to_le_bytes();
804        buf[len - 4..len].copy_from_slice(&bytes);
805        self.0 = &mut buf[..len - 4];
806        Ok(())
807    }
808}
809
810impl<'a> ForwardReader for &'a [u8] {
811    #[inline]
812    fn read_byte(&mut self) -> Option<u8> {
813        if self.is_empty() {
814            return None;
815        }
816        let b = self[0];
817        *self = &self[1..];
818        Some(b)
819    }
820
821    #[inline]
822    fn read_u32_le(&mut self) -> Option<u32> {
823        if self.len() < 4 {
824            return None;
825        }
826        let v = u32::from_le_bytes([self[0], self[1], self[2], self[3]]);
827        *self = &self[4..];
828        Some(v)
829    }
830}
831
832// ---------------------------------------------------------------------------
833// Two-state interleaved byte rANS
834// ---------------------------------------------------------------------------
835
836/// Two-state interleaved byte rANS encoder.
837///
838/// Encodes symbols into two interleaved streams (rans0, rans1).
839/// Equivalent to the two-stream interleaving pattern in `main.cpp`.
840pub struct ByteInterleavedEncoder<'a, W: BackwardWriter> {
841    state0: RansByteState,
842    state1: RansByteState,
843    writer: &'a mut W,
844    _scale_bits: u32,
845    _num_symbols: usize,
846    _odd_symbol: Option<u8>,
847}
848
849impl<'a, W: BackwardWriter> ByteInterleavedEncoder<'a, W> {
850    /// Create a new interleaved encoder.
851    pub fn new(writer: &'a mut W, scale_bits: u32) -> Self {
852        Self {
853            state0: RansByteState::new(),
854            state1: RansByteState::new(),
855            writer,
856            _scale_bits: scale_bits,
857            _num_symbols: 0,
858            _odd_symbol: None,
859        }
860    }
861
862    /// Encode symbols in reverse order (last to first).
863    ///
864    /// Pass the symbol array slice in forward order; this function iterates
865    /// it in reverse internally.
866    pub fn encode_reverse(
867        &mut self,
868        symbols: &[u8],
869        esyms: &[RansByteEncSymbol],
870    ) -> Result<(), EncodeError> {
871        let n = symbols.len();
872        self._num_symbols = n;
873
874        if n == 0 {
875            return Ok(());
876        }
877
878        // Handle odd length: last symbol goes to state0
879        if n & 1 != 0 {
880            let s = symbols[n - 1];
881            rans_byte_enc_put_symbol(&mut self.state0, &mut *self.writer, &esyms[s as usize])?;
882        }
883
884        // Process pairs in reverse
885        let mut i = n & !1;
886        while i > 0 {
887            let s1 = symbols[i - 1] as usize;
888            let s0 = symbols[i - 2] as usize;
889
890            // Interleave: first encode state1, then state0
891            // (reverse of decoding order)
892            rans_byte_enc_put_symbol(&mut self.state1, &mut *self.writer, &esyms[s1])?;
893            rans_byte_enc_put_symbol(&mut self.state0, &mut *self.writer, &esyms[s0])?;
894
895            i = i.wrapping_sub(2);
896        }
897
898        Ok(())
899    }
900
901    /// Flush both encoder states (state1 first, then state0).
902    ///
903    /// Matches the flush order in the upstream interleaving example.
904    pub fn flush(&mut self) -> Result<(), EncodeError> {
905        rans_byte_enc_flush(&self.state1, &mut *self.writer)?;
906        rans_byte_enc_flush(&self.state0, &mut *self.writer)?;
907        Ok(())
908    }
909
910    /// Finalize: encode, flush, return the encoded slice.
911    pub fn finalize(
912        mut self,
913        symbols: &[u8],
914        esyms: &[RansByteEncSymbol],
915    ) -> Result<(), EncodeError> {
916        self.encode_reverse(symbols, esyms)?;
917        self.flush()
918    }
919}
920
921/// Two-state interleaved byte rANS decoder.
922pub struct ByteInterleavedDecoder<'a, R: ForwardReader> {
923    state0: RansByteState,
924    state1: RansByteState,
925    reader: &'a mut R,
926    scale_bits: u32,
927}
928
929impl<'a, R: ForwardReader> ByteInterleavedDecoder<'a, R> {
930    /// Create a new interleaved decoder.
931    pub fn new(reader: &'a mut R, scale_bits: u32) -> Result<Self, DecodeError> {
932        let state0 = rans_byte_dec_init(&mut *reader)?;
933        let state1 = rans_byte_dec_init(&mut *reader)?;
934        Ok(Self {
935            state0,
936            state1,
937            reader,
938            scale_bits,
939        })
940    }
941
942    /// Decode all symbols into the output buffer.
943    ///
944    /// Returns the number of symbols decoded.
945    pub fn decode(
946        &mut self,
947        output: &mut [u8],
948        cum2sym: &[u8],
949        dsyms: &[RansByteDecSymbol],
950    ) -> Result<usize, DecodeError> {
951        let n = output.len();
952        if n == 0 {
953            return Ok(0);
954        }
955
956        let even_n = n & !1;
957
958        // Process pairs
959        let mut i = 0usize;
960        while i < even_n {
961            let cf0 = rans_byte_dec_get(&self.state0, self.scale_bits);
962            let s0 = cum2sym[cf0 as usize] as usize;
963            let cf1 = rans_byte_dec_get(&self.state1, self.scale_bits);
964            let s1 = cum2sym[cf1 as usize] as usize;
965
966            output[i] = s0 as u8;
967            output[i + 1] = s1 as u8;
968
969            rans_byte_dec_advance_symbol_step(&mut self.state0, &dsyms[s0], self.scale_bits);
970            rans_byte_dec_advance_symbol_step(&mut self.state1, &dsyms[s1], self.scale_bits);
971            rans_byte_dec_renorm(&mut self.state0, &mut *self.reader)?;
972            rans_byte_dec_renorm(&mut self.state1, &mut *self.reader)?;
973
974            i += 2;
975        }
976
977        // Handle odd byte
978        if n & 1 != 0 {
979            let cf0 = rans_byte_dec_get(&self.state0, self.scale_bits);
980            let s0 = cum2sym[cf0 as usize] as usize;
981            output[n - 1] = s0 as u8;
982            rans_byte_dec_advance_symbol(
983                &mut self.state0,
984                &mut *self.reader,
985                &dsyms[s0],
986                self.scale_bits,
987            )?;
988        }
989
990        Ok(n)
991    }
992
993    /// Get the final decoder states (for verification).
994    pub fn states(&self) -> (RansByteState, RansByteState) {
995        (self.state0, self.state1)
996    }
997}
998
999// ---------------------------------------------------------------------------
1000// 64-bit word-aligned rANS (rans64.h reconstruction)
1001// ---------------------------------------------------------------------------
1002
1003/// Lower bound of the normalization interval for 64-bit rANS.
1004///
1005/// Equivalent to `RANS64_L` in `rans64.h`.
1006pub const RANS64_L: u64 = 1u64 << 31;
1007
1008/// 64-bit rANS encoder/decoder state.
1009///
1010/// Equivalent to `Rans64State` (typedef for `uint64_t`) in `rans64.h`.
1011/// Uses 63 bits of effective state space.
1012#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1013pub struct Rans64State(pub u64);
1014
1015impl Rans64State {
1016    /// Create a new state initialized to the lower bound.
1017    #[inline]
1018    pub const fn new() -> Self {
1019        Self(RANS64_L)
1020    }
1021
1022    /// Return the raw state value.
1023    #[inline]
1024    pub const fn get(&self) -> u64 {
1025        self.0
1026    }
1027}
1028
1029impl Default for Rans64State {
1030    #[inline]
1031    fn default() -> Self {
1032        Self::new()
1033    }
1034}
1035
1036// ---------------------------------------------------------------------------
1037// Backward word-32 writer
1038// ---------------------------------------------------------------------------
1039
1040/// A backward-growing writer that writes 32-bit words.
1041///
1042/// Starts at the logical end of the provided byte buffer and writes u32
1043/// values (little-endian, 4 bytes each) toward the beginning.
1044/// Used for 64-bit rANS encoder output.
1045///
1046/// Equivalent to `uint32_t* ptr` pointing past the end of the output buffer
1047/// in the upstream `rans64.h`.
1048pub struct BackwardWord32Writer<'a> {
1049    buf: &'a mut [u8],
1050    pos: usize, // current write position (bytes, 0 <= pos <= len, always multiple of 4)
1051}
1052
1053impl<'a> BackwardWord32Writer<'a> {
1054    /// Create a new backward word-32 writer from a mutable byte slice.
1055    ///
1056    /// The writer starts at the end of the buffer.
1057    #[inline]
1058    pub fn new(buf: &'a mut [u8]) -> Self {
1059        let len = buf.len();
1060        debug_assert!(
1061            len % 4 == 0,
1062            "BackwardWord32Writer buffer len must be multiple of 4"
1063        );
1064        Self { buf, pos: len }
1065    }
1066
1067    /// Write a single u32 at the current position (decrementing the cursor by 4).
1068    ///
1069    /// Returns `Ok(())` if there was room, `Err(())` if the writer has
1070    /// exhausted its buffer.
1071    #[inline]
1072    pub fn write_word32(&mut self, v: u32) -> Result<(), ()> {
1073        if self.pos < 4 {
1074            return Err(());
1075        }
1076        self.pos -= 4;
1077        self.buf[self.pos..self.pos + 4].copy_from_slice(&v.to_le_bytes());
1078        Ok(())
1079    }
1080
1081    /// Current zero-based write position (bytes before the cursor).
1082    #[inline]
1083    pub fn position(&self) -> usize {
1084        self.pos
1085    }
1086
1087    /// Number of bytes written so far.
1088    #[inline]
1089    pub fn bytes_written(&self) -> usize {
1090        self.buf.len() - self.pos
1091    }
1092
1093    /// Number of u32 words written so far.
1094    #[inline]
1095    pub fn words_written(&self) -> usize {
1096        (self.buf.len() - self.pos) / 4
1097    }
1098
1099    /// Return the encoded portion (from current position to end) as a slice.
1100    #[inline]
1101    pub fn encoded(&self) -> &[u8] {
1102        &self.buf[self.pos..]
1103    }
1104
1105    /// Return the remaining capacity (number of bytes that can still be written).
1106    #[inline]
1107    pub fn remaining(&self) -> usize {
1108        self.pos
1109    }
1110}
1111
1112impl fmt::Debug for BackwardWord32Writer<'_> {
1113    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1114        f.debug_struct("BackwardWord32Writer")
1115            .field("pos", &self.pos)
1116            .field("len", &self.buf.len())
1117            .finish()
1118    }
1119}
1120
1121// ---------------------------------------------------------------------------
1122// Forward word-32 reader
1123// ---------------------------------------------------------------------------
1124
1125/// A forward-growing reader that reads 32-bit words.
1126///
1127/// Used for 64-bit rANS decoding input. Reads u32 values (little-endian,
1128/// 4 bytes each) from the beginning of the compressed stream moving forward.
1129///
1130/// Equivalent to `uint32_t* ptr` pointing to the start of compressed data
1131/// in the upstream `rans64.h`.
1132#[derive(Clone)]
1133pub struct Word32Reader<'a> {
1134    buf: &'a [u8],
1135    pos: usize,
1136}
1137
1138impl<'a> Word32Reader<'a> {
1139    /// Create a new word-32 reader from a byte slice.
1140    #[inline]
1141    pub fn new(buf: &'a [u8]) -> Self {
1142        Self { buf, pos: 0 }
1143    }
1144
1145    /// Read a single u32 (little-endian) at the current position, advancing by 4.
1146    ///
1147    /// Returns `None` if fewer than 4 bytes remain.
1148    #[inline]
1149    pub fn read_word32(&mut self) -> Option<u32> {
1150        if self.pos + 4 > self.buf.len() {
1151            return None;
1152        }
1153        let v = u32::from_le_bytes([
1154            self.buf[self.pos],
1155            self.buf[self.pos + 1],
1156            self.buf[self.pos + 2],
1157            self.buf[self.pos + 3],
1158        ]);
1159        self.pos += 4;
1160        Some(v)
1161    }
1162
1163    /// Current read position in bytes.
1164    #[inline]
1165    pub fn position(&self) -> usize {
1166        self.pos
1167    }
1168
1169    /// Number of bytes consumed so far.
1170    #[inline]
1171    pub fn bytes_consumed(&self) -> usize {
1172        self.pos
1173    }
1174
1175    /// Number of u32 words consumed so far.
1176    #[inline]
1177    pub fn words_consumed(&self) -> usize {
1178        self.pos / 4
1179    }
1180
1181    /// Remaining unread bytes.
1182    #[inline]
1183    pub fn remaining(&self) -> usize {
1184        self.buf.len().saturating_sub(self.pos)
1185    }
1186}
1187
1188impl fmt::Debug for Word32Reader<'_> {
1189    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1190        f.debug_struct("Word32Reader")
1191            .field("pos", &self.pos)
1192            .field("len", &self.buf.len())
1193            .finish()
1194    }
1195}
1196
1197// ---------------------------------------------------------------------------
1198// 64-bit encoder symbol (reciprocal-multiply setup)
1199// ---------------------------------------------------------------------------
1200
1201/// Precomputed reciprocal encoder symbol for 64-bit rANS.
1202///
1203/// Equivalent to `Rans64EncSymbol` in `rans64.h`.
1204///
1205/// The fast encoder uses a 64x64 multiply-high approximation to divide
1206/// by the symbol frequency, avoiding an integer division in the hot loop.
1207/// The reciprocal setup requires 128-bit arithmetic.
1208#[derive(Debug, Clone, Copy)]
1209pub struct Rans64EncSymbol {
1210    /// (Exclusive) upper bound of pre-normalization interval.
1211    pub x_max: u64,
1212    /// Fixed-point reciprocal frequency (u64).
1213    pub rcp_freq: u64,
1214    /// Bias value.
1215    pub bias: u64,
1216    /// Complement of frequency: `(1 << scale_bits) - freq` (32-bit, not 16-bit!)
1217    pub cmpl_freq: u32,
1218    /// Reciprocal shift amount.
1219    pub rcp_shift: u32,
1220}
1221
1222impl Rans64EncSymbol {
1223    /// Initialize a 64-bit encoder symbol with validation.
1224    ///
1225    /// Equivalent to `Rans64EncSymbolInit` in `rans64.h`.
1226    ///
1227    /// Calculates the 64-bit reciprocal using 128-bit division:
1228    /// `rcp_freq = ((1 << (shift + 63)) + freq - 1) / freq`.
1229    ///
1230    /// Returns `Err(ModelError::InvalidScaleBits)` if `scale_bits` is not in
1231    /// `1..=31`. Returns `Err(ModelError::ZeroFrequency)` if `freq == 0`.
1232    /// Returns `Err(ModelError::StartOutOfRange)` if `start > (1 << scale_bits)`.
1233    /// Returns `Err(ModelError::FrequencyOutOfRange)` if
1234    /// `freq > (1 << scale_bits) - start`.
1235    #[inline]
1236    pub fn new(start: u32, freq: u32, scale_bits: u32) -> Result<Self, ModelError> {
1237        if !(1..=31).contains(&scale_bits) {
1238            return Err(ModelError::InvalidScaleBits);
1239        }
1240        let max_start = 1u64 << scale_bits;
1241        if (start as u64) > max_start {
1242            return Err(ModelError::StartOutOfRange);
1243        }
1244        if freq == 0 {
1245            return Err(ModelError::ZeroFrequency);
1246        }
1247        if (freq as u64) > max_start - (start as u64) {
1248            return Err(ModelError::FrequencyOutOfRange);
1249        }
1250
1251        Ok(Self::new_unchecked(start, freq, scale_bits))
1252    }
1253
1254    /// Initialize a 64-bit encoder symbol without validation.
1255    ///
1256    /// Caller must ensure preconditions are met. Prefer [`new`] for
1257    /// external use.
1258    #[inline]
1259    pub(crate) fn new_unchecked(start: u32, freq: u32, scale_bits: u32) -> Self {
1260        debug_assert!(scale_bits <= 31, "scale_bits must be <= 31");
1261        debug_assert!((start as u64) <= (1u64 << scale_bits), "start out of range");
1262        debug_assert!(
1263            (freq as u64) <= (1u64 << scale_bits) - (start as u64),
1264            "freq out of range"
1265        );
1266
1267        // x_max = ((RANS64_L >> scale_bits) << 32) * freq
1268        let x_max = ((RANS64_L >> scale_bits) << 32) * (freq as u64);
1269        // cmpl_freq is u32 because scale_bits may reach 31, giving a complement
1270        // of up to 2^31 - 1, which does not fit in u16.
1271        let cmpl_freq = ((1u64 << scale_bits) - freq as u64) as u32;
1272
1273        if freq < 2 {
1274            // freq == 1 special case
1275            // rcp_freq = ~0u64, rcp_shift = 0
1276            // bias = start + (1 << scale_bits) - 1
1277            Self {
1278                x_max,
1279                rcp_freq: !0u64,
1280                rcp_shift: 0,
1281                bias: (start as u64) + (1u64 << scale_bits) - 1,
1282                cmpl_freq: (1u64 << scale_bits) as u32 - freq,
1283            }
1284        } else {
1285            // Alverson "Integer Division using Reciprocals"
1286            // Find smallest shift such that freq <= (1 << shift)
1287            let mut shift = 0u32;
1288            while freq > (1u32 << shift) {
1289                shift += 1;
1290            }
1291
1292            // rcp_freq = ((1u128 << (shift + 63)) + freq - 1) / freq
1293            // This is a 128-bit numerator, which we compute with u128 to avoid
1294            // overflow and allocate the result into u64.
1295            let rcp_freq = (((1u128 << (shift + 63)) + (freq as u128) - 1) / (freq as u128)) as u64;
1296            let rcp_shift = shift - 1;
1297
1298            Self {
1299                x_max,
1300                rcp_freq,
1301                rcp_shift,
1302                bias: start as u64,
1303                cmpl_freq,
1304            }
1305        }
1306    }
1307}
1308
1309// ---------------------------------------------------------------------------
1310// 64-bit decoder symbol
1311// ---------------------------------------------------------------------------
1312
1313/// Decoder symbol description for 64-bit rANS.
1314///
1315/// Equivalent to `Rans64DecSymbol` in `rans64.h`.
1316/// start and freq are u32 (scale_bits up to 31).
1317#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1318pub struct Rans64DecSymbol {
1319    /// Start of range.
1320    pub start: u32,
1321    /// Symbol frequency.
1322    pub freq: u32,
1323}
1324
1325impl Rans64DecSymbol {
1326    /// Initialize a 64-bit decoder symbol with validation.
1327    ///
1328    /// Equivalent to `Rans64DecSymbolInit`.
1329    ///
1330    /// Returns `Err(ModelError::ZeroFrequency)` if `freq == 0`.
1331    /// Returns `Err(ModelError::StartOutOfRange)` if `start > (1 << 31)`.
1332    /// Returns `Err(ModelError::FrequencyOutOfRange)` if
1333    /// `freq > (1 << 31) - start`.
1334    #[inline]
1335    pub fn new(start: u32, freq: u32) -> Result<Self, ModelError> {
1336        if freq == 0 {
1337            return Err(ModelError::ZeroFrequency);
1338        }
1339        if (start as u64) > (1u64 << 31) {
1340            return Err(ModelError::StartOutOfRange);
1341        }
1342        if (freq as u64) > (1u64 << 31) - (start as u64) {
1343            return Err(ModelError::FrequencyOutOfRange);
1344        }
1345        Ok(Self::new_unchecked(start, freq))
1346    }
1347
1348    /// Initialize a 64-bit decoder symbol without validation.
1349    ///
1350    /// Caller must ensure preconditions are met. Prefer [`new`] for
1351    /// external use.
1352    #[inline]
1353    pub(crate) fn new_unchecked(start: u32, freq: u32) -> Self {
1354        debug_assert!((start as u64) <= (1u64 << 31), "start out of range");
1355        debug_assert!(
1356            (freq as u64) <= (1u64 << 31) - (start as u64),
1357            "freq out of range"
1358        );
1359        Self { start, freq }
1360    }
1361}
1362
1363// ---------------------------------------------------------------------------
1364// 64-bit multiply-high helper
1365// ---------------------------------------------------------------------------
1366
1367/// Compute the high 64 bits of the 128-bit product `a * b`.
1368///
1369/// Equivalent to `Rans64MulHi` in `rans64.h`:
1370/// `((unsigned __int128)a * b) >> 64`.
1371#[inline]
1372pub fn rans64_mul_hi(a: u64, b: u64) -> u64 {
1373    ((a as u128) * (b as u128) >> 64) as u64
1374}
1375
1376// ---------------------------------------------------------------------------
1377// Division-based encoder (reference path)
1378// ---------------------------------------------------------------------------
1379
1380/// Renormalize the 64-bit encoder state by emitting u32 words (LSB first).
1381///
1382/// Equivalent to `Rans64EncRenorm` in `rans64.h`.
1383#[inline]
1384pub fn rans64_enc_renorm(
1385    x: u64,
1386    writer: &mut BackwardWord32Writer,
1387    freq: u32,
1388    scale_bits: u32,
1389) -> Result<u64, EncodeError> {
1390    let x_max = ((RANS64_L >> scale_bits) << 32) * (freq as u64);
1391    let mut x = x;
1392    if x >= x_max {
1393        while x >= x_max {
1394            writer
1395                .write_word32((x & 0xffffffff) as u32)
1396                .map_err(|_| EncodeError::OutputTooSmall)?;
1397            x >>= 32;
1398        }
1399    }
1400    Ok(x)
1401}
1402
1403/// 64-bit encoder put-symbol (division-based reference path).
1404///
1405/// Equivalent to `Rans64EncPut` in `rans64.h`.
1406#[inline]
1407pub fn rans64_enc_put(
1408    state: &mut Rans64State,
1409    writer: &mut BackwardWord32Writer,
1410    start: u32,
1411    freq: u32,
1412    scale_bits: u32,
1413) -> Result<(), EncodeError> {
1414    let x = rans64_enc_renorm(state.0, writer, freq, scale_bits)?;
1415    // C(s, x) = ((x / freq) << scale_bits) + (x % freq) + start
1416    state.0 = ((x / (freq as u64)) << scale_bits) + (x % (freq as u64)) + (start as u64);
1417    Ok(())
1418}
1419
1420/// Flush the 64-bit encoder.
1421///
1422/// Equivalent to `Rans64EncFlush` in `rans64.h`.
1423///
1424/// Writes the full 64-bit state as two u32 words in little-endian order.
1425/// Low word first, then high word (decrementing the pointer).
1426#[inline]
1427pub fn rans64_enc_flush(
1428    state: &Rans64State,
1429    writer: &mut BackwardWord32Writer,
1430) -> Result<(), EncodeError> {
1431    let x = state.0;
1432    // Write low word, then high word (both move backward)
1433    writer
1434        .write_word32((x >> 32) as u32)
1435        .map_err(|_| EncodeError::OutputTooSmall)?;
1436    writer
1437        .write_word32((x & 0xffffffff) as u32)
1438        .map_err(|_| EncodeError::OutputTooSmall)?;
1439    Ok(())
1440}
1441
1442// ---------------------------------------------------------------------------
1443// Division-based decoder (reference path)
1444// ---------------------------------------------------------------------------
1445
1446/// Initialize the 64-bit decoder.
1447///
1448/// Equivalent to `Rans64DecInit` in `rans64.h`.
1449///
1450/// Reads two u32 words (low word first, then high word) to reconstruct
1451/// the initial 64-bit state.
1452#[inline]
1453pub fn rans64_dec_init(reader: &mut Word32Reader) -> Result<Rans64State, DecodeError> {
1454    let lo = reader.read_word32().ok_or(DecodeError::InputTooShort)?;
1455    let hi = reader.read_word32().ok_or(DecodeError::InputTooShort)?;
1456    Ok(Rans64State((lo as u64) | ((hi as u64) << 32)))
1457}
1458
1459/// Get the cumulative frequency from the current 64-bit state.
1460///
1461/// Equivalent to `Rans64DecGet` in `rans64.h`.
1462#[inline]
1463pub fn rans64_dec_get(state: &Rans64State, scale_bits: u32) -> u32 {
1464    (state.0 & ((1u64 << scale_bits) - 1)) as u32
1465}
1466
1467/// Advance the 64-bit decoder by a symbol (with renormalization).
1468///
1469/// Equivalent to `Rans64DecAdvance` in `rans64.h`.
1470#[inline]
1471pub fn rans64_dec_advance(
1472    state: &mut Rans64State,
1473    reader: &mut Word32Reader,
1474    start: u32,
1475    freq: u32,
1476    scale_bits: u32,
1477) -> Result<(), DecodeError> {
1478    let mask = (1u64 << scale_bits) - 1;
1479    let x = state.0;
1480    let mut x = (freq as u64) * (x >> scale_bits) + (x & mask) - (start as u64);
1481
1482    // renormalize: read u32 words until x >= RANS64_L
1483    if x < RANS64_L {
1484        loop {
1485            let word = reader.read_word32().ok_or(DecodeError::InputTooShort)?;
1486            x = (x << 32) | (word as u64);
1487            if x >= RANS64_L {
1488                break;
1489            }
1490        }
1491    }
1492
1493    state.0 = x;
1494    Ok(())
1495}
1496
1497// ---------------------------------------------------------------------------
1498// Reciprocal-multiply encoder (fast path)
1499// ---------------------------------------------------------------------------
1500
1501/// Encoder put-symbol using precomputed reciprocal (fast path).
1502///
1503/// Equivalent to `Rans64EncPutSymbol` in `rans64.h`.
1504///
1505/// Uses the 64-bit multiply-high approximation (`rans64_mul_hi`) to avoid
1506/// integer division in the hot loop.
1507#[inline]
1508pub fn rans64_enc_put_symbol(
1509    state: &mut Rans64State,
1510    writer: &mut BackwardWord32Writer,
1511    sym: &Rans64EncSymbol,
1512) -> Result<(), EncodeError> {
1513    debug_assert!(sym.x_max != 0, "cannot encode symbol with freq=0");
1514
1515    let mut x = state.0;
1516    if x >= sym.x_max {
1517        while x >= sym.x_max {
1518            writer
1519                .write_word32((x & 0xffffffff) as u32)
1520                .map_err(|_| EncodeError::OutputTooSmall)?;
1521            x >>= 32;
1522        }
1523    }
1524
1525    // x = C(s, x) using reciprocal multiply
1526    // q = Rans64MulHi(x, rcp_freq) >> rcp_shift
1527    let q = rans64_mul_hi(x, sym.rcp_freq) >> sym.rcp_shift;
1528    state.0 = x + sym.bias + q * (sym.cmpl_freq as u64);
1529    Ok(())
1530}
1531
1532// ---------------------------------------------------------------------------
1533// Decoder advance with symbol (convenience)
1534// ---------------------------------------------------------------------------
1535
1536/// Advance the 64-bit decoder using a decoder symbol.
1537///
1538/// Equivalent to `Rans64DecAdvanceSymbol` in `rans64.h`.
1539#[inline]
1540pub fn rans64_dec_advance_symbol(
1541    state: &mut Rans64State,
1542    reader: &mut Word32Reader,
1543    sym: &Rans64DecSymbol,
1544    scale_bits: u32,
1545) -> Result<(), DecodeError> {
1546    rans64_dec_advance(state, reader, sym.start, sym.freq, scale_bits)
1547}
1548
1549// ---------------------------------------------------------------------------
1550// Step-only operations (for interleaving)
1551// ---------------------------------------------------------------------------
1552
1553/// Advance the 64-bit decoder without renormalization (interleaving step).
1554///
1555/// Equivalent to `Rans64DecAdvanceStep` in `rans64.h`.
1556#[inline]
1557pub fn rans64_dec_advance_step(state: &mut Rans64State, start: u32, freq: u32, scale_bits: u32) {
1558    let mask = (1u64 << scale_bits) - 1;
1559    let x = state.0;
1560    state.0 = (freq as u64) * (x >> scale_bits) + (x & mask) - (start as u64);
1561}
1562
1563/// Advance the 64-bit decoder step using a decoder symbol.
1564///
1565/// Equivalent to `Rans64DecAdvanceSymbolStep`.
1566#[inline]
1567pub fn rans64_dec_advance_symbol_step(
1568    state: &mut Rans64State,
1569    sym: &Rans64DecSymbol,
1570    scale_bits: u32,
1571) {
1572    rans64_dec_advance_step(state, sym.start, sym.freq, scale_bits);
1573}
1574
1575/// 64-bit decoder renormalization only.
1576///
1577/// Equivalent to `Rans64DecRenorm` in `rans64.h`.
1578#[inline]
1579pub fn rans64_dec_renorm(
1580    state: &mut Rans64State,
1581    reader: &mut Word32Reader,
1582) -> Result<(), DecodeError> {
1583    let mut x = state.0;
1584    if x < RANS64_L {
1585        loop {
1586            let word = reader.read_word32().ok_or(DecodeError::InputTooShort)?;
1587            x = (x << 32) | (word as u64);
1588            if x >= RANS64_L {
1589                break;
1590            }
1591        }
1592        state.0 = x;
1593    }
1594    Ok(())
1595}
1596
1597// ---------------------------------------------------------------------------
1598// Word-aligned rANS (word rANS — rans_word_sse41.h scalar path)
1599// ---------------------------------------------------------------------------
1600
1601/// Lower bound of the normalization interval for word rANS.
1602///
1603/// Equivalent to `RANS_WORD_L` in `rans_word_sse41.h`.
1604/// Word rANS uses L=2^16 with 16-bit renormalization words.
1605/// Default scale_bits for word rANS (per upstream design).
1606/// Upstream rans_word_sse41.h hardcodes this to 12.
1607/// The table has 1<<12 = 4096 slots, and all encode/decode
1608/// formulas assume this value.
1609pub const RANS_WORD_SCALE_BITS: u32 = 12;
1610
1611/// Lower bound of the normalization interval for word rANS.
1612///
1613/// Equivalent to `RANS_WORD_L` in `rans_word_sse41.h`.
1614/// Word rANS uses L=2^16 with 16-bit renormalization words.
1615pub const RANS_WORD_L: u32 = 1u32 << 16;
1616
1617/// Word rANS encoder/decoder state.
1618///
1619/// Equivalent to `RansWordEnc` / `RansWordDec` (typedef for `uint32_t`).
1620/// Uses 32-bit state space with L=2^16 and 16-bit word renormalization.
1621#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1622pub struct RansWordState(pub u32);
1623
1624impl RansWordState {
1625    /// Create a new state initialized to the lower bound.
1626    #[inline]
1627    pub const fn new() -> Self {
1628        Self(RANS_WORD_L)
1629    }
1630
1631    /// Return the raw state value.
1632    #[inline]
1633    pub const fn get(&self) -> u32 {
1634        self.0
1635    }
1636}
1637
1638/// Backward writer for 16-bit words (word rANS encoder output).
1639///
1640/// Writes u16 values in little-endian order to a byte buffer,
1641/// starting from the end and moving toward the beginning.
1642#[derive(Debug)]
1643pub struct BackwardWord16Writer<'a> {
1644    buf: &'a mut [u8],
1645    pos: usize,
1646}
1647
1648impl<'a> BackwardWord16Writer<'a> {
1649    /// Create a new backward word writer from a mutable byte slice.
1650    ///
1651    /// The writer starts at the end of the buffer.
1652    #[inline]
1653    pub fn new(buf: &'a mut [u8]) -> Self {
1654        let len = buf.len();
1655        Self { buf, pos: len }
1656    }
1657
1658    /// Write a single u16 word in little-endian format.
1659    ///
1660    /// Returns `Ok(())` if there was room, `Err(())` if the buffer is exhausted.
1661    #[inline]
1662    pub fn write_word16(&mut self, v: u16) -> Result<(), ()> {
1663        if self.pos < 2 {
1664            return Err(());
1665        }
1666        self.pos -= 2;
1667        self.buf[self.pos..self.pos + 2].copy_from_slice(&v.to_le_bytes());
1668        Ok(())
1669    }
1670
1671    /// Current write position (0 = buffer exhausted).
1672    #[inline]
1673    pub fn position(&self) -> usize {
1674        self.pos
1675    }
1676
1677    /// Return the encoded portion (from current position to end) as a slice.
1678    #[inline]
1679    pub fn encoded(&self) -> &[u8] {
1680        &self.buf[self.pos..]
1681    }
1682
1683    /// Number of bytes written so far.
1684    #[inline]
1685    pub fn bytes_written(&self) -> usize {
1686        self.buf.len() - self.pos
1687    }
1688}
1689
1690/// Forward reader for 16-bit words (word rANS decoder input).
1691///
1692/// Reads u16 values in little-endian order from a byte slice,
1693/// starting from the beginning and moving forward.
1694#[derive(Debug, Clone)]
1695pub struct Word16Reader<'a> {
1696    buf: &'a [u8],
1697    pos: usize,
1698}
1699
1700impl<'a> Word16Reader<'a> {
1701    /// Create a new word reader from a byte slice.
1702    #[inline]
1703    pub fn new(buf: &'a [u8]) -> Self {
1704        Self { buf, pos: 0 }
1705    }
1706
1707    /// Read a single u16 word in little-endian format.
1708    ///
1709    /// Returns `None` if fewer than 2 bytes remain.
1710    #[inline]
1711    pub fn read_word16(&mut self) -> Option<u16> {
1712        if self.pos + 2 > self.buf.len() {
1713            return None;
1714        }
1715        let v = u16::from_le_bytes([self.buf[self.pos], self.buf[self.pos + 1]]);
1716        self.pos += 2;
1717        Some(v)
1718    }
1719
1720    /// Current read position (bytes consumed so far).
1721    #[inline]
1722    pub fn position(&self) -> usize {
1723        self.pos
1724    }
1725
1726    /// Number of bytes consumed.
1727    #[inline]
1728    pub fn bytes_consumed(&self) -> usize {
1729        self.pos
1730    }
1731}
1732
1733// ---------------------------------------------------------------------------
1734// Word rANS table types
1735// ---------------------------------------------------------------------------
1736
1737/// A single slot in the word rANS decode table.
1738///
1739/// Equivalent to `RansWordSlot` in `rans_word_sse41.h`.
1740#[derive(Debug, Clone, Copy)]
1741pub struct RansWordSlot {
1742    /// Symbol frequency.
1743    pub freq: u16,
1744    /// Bias (position within the symbol's frequency range).
1745    pub bias: u16,
1746}
1747
1748/// Word rANS decode tables.
1749///
1750/// Equivalent to `RansWordTables` in `rans_word_sse41.h`.
1751/// Contains `M` slots (one per cumulative frequency) and a
1752/// slot-to-symbol mapping.
1753pub struct RansWordTables<'a> {
1754    pub slots: &'a [RansWordSlot],
1755    pub slot2sym: &'a [u8],
1756}
1757
1758/// Check that scale_bits matches the word rANS upstream constant.
1759/// Returns `Err(ModelError::InvalidScaleBits)` for unsupported values.
1760#[inline]
1761pub fn rans_word_check_scale_bits(scale_bits: u32) -> Result<(), ModelError> {
1762    if scale_bits != RANS_WORD_SCALE_BITS {
1763        return Err(ModelError::InvalidScaleBits);
1764    }
1765    Ok(())
1766}
1767
1768// ---------------------------------------------------------------------------
1769// Word rANS encoder functions
1770// ---------------------------------------------------------------------------
1771
1772/// Initialize a word rANS encoder state.
1773///
1774/// Equivalent to `RansWordEncInit` in `rans_word_sse41.h`.
1775#[inline]
1776pub fn rans_word_enc_init() -> RansWordState {
1777    RansWordState::new()
1778}
1779
1780/// Word rANS encoder renormalization.
1781///
1782/// Emits the low 16 bits when the state exceeds the threshold.
1783/// Returns `Err(EncodeError::OutputTooSmall)` if the buffer is exhausted.
1784#[inline]
1785pub fn rans_word_enc_renorm(
1786    state: &mut RansWordState,
1787    writer: &mut BackwardWord16Writer,
1788    freq: u32,
1789    scale_bits: u32,
1790) -> Result<(), EncodeError> {
1791    let x = state.0;
1792    // Threshold: (L >> scale_bits) << 16 * freq
1793    // For scale_bits=12: (0x10000 >> 12) << 16 * freq = (16) << 16 * freq = 0x100000 * freq
1794    let threshold = ((RANS_WORD_L >> scale_bits) << 16) * freq;
1795    if x >= threshold {
1796        writer
1797            .write_word16((x & 0xffff) as u16)
1798            .map_err(|_| EncodeError::OutputTooSmall)?;
1799        state.0 = x >> 16;
1800    }
1801    Ok(())
1802}
1803
1804/// Word rANS encoder — encode a single symbol using division.
1805///
1806/// Equivalent to `RansWordEncPut` in `rans_word_sse41.h`.
1807/// Formula: x = C(s,x) = ((x/freq) << scale_bits) + (x%freq) + start
1808#[inline]
1809pub fn rans_word_enc_put(
1810    state: &mut RansWordState,
1811    writer: &mut BackwardWord16Writer,
1812    start: u32,
1813    freq: u32,
1814    scale_bits: u32,
1815) -> Result<(), EncodeError> {
1816    // Validate scale_bits matches upstream constant
1817    rans_word_check_scale_bits(scale_bits).map_err(|_| EncodeError::OutputTooSmall)?;
1818
1819    // Renormalize
1820    rans_word_enc_renorm(state, writer, freq, scale_bits)?;
1821
1822    // x = C(s,x)
1823    let x = state.0;
1824    state.0 = ((x / freq) << scale_bits) + (x % freq) + start;
1825    Ok(())
1826}
1827
1828/// Flush a word rANS encoder — write the final state to the buffer.
1829///
1830/// Equivalent to `RansWordEncFlush` in `rans_word_sse41.h`.
1831/// Writes 2 u16 words (4 bytes total).
1832#[inline]
1833pub fn rans_word_enc_flush(
1834    state: &RansWordState,
1835    writer: &mut BackwardWord16Writer,
1836) -> Result<(), EncodeError> {
1837    let x = state.0;
1838    writer
1839        .write_word16((x >> 16) as u16)
1840        .map_err(|_| EncodeError::OutputTooSmall)?;
1841    writer
1842        .write_word16((x & 0xffff) as u16)
1843        .map_err(|_| EncodeError::OutputTooSmall)?;
1844    Ok(())
1845}
1846
1847// ---------------------------------------------------------------------------
1848// Word rANS decoder functions
1849// ---------------------------------------------------------------------------
1850
1851/// Initialize a word rANS decoder from a compressed stream.
1852///
1853/// Equivalent to `RansWordDecInit` in `rans_word_sse41.h`.
1854/// Reads 2 u16 words (4 bytes) from the stream.
1855#[inline]
1856pub fn rans_word_dec_init(reader: &mut Word16Reader) -> Result<RansWordState, DecodeError> {
1857    let lo = reader.read_word16().ok_or(DecodeError::InputTooShort)? as u32;
1858    let hi = reader.read_word16().ok_or(DecodeError::InputTooShort)? as u32;
1859    Ok(RansWordState(lo | (hi << 16)))
1860}
1861
1862/// Word rANS decoder — decode one symbol and advance the state.
1863///
1864/// Equivalent to `RansWordDecSym` in `rans_word_sse41.h`.
1865/// Uses the table-based decoder: slot = x & (M-1), then
1866/// x' = slots[slot].freq * (x >> scale_bits) + slots[slot].bias
1867/// Returns the decoded symbol.
1868#[inline]
1869pub fn rans_word_dec_sym(
1870    state: &mut RansWordState,
1871    tables: &RansWordTables<'_>,
1872    scale_bits: u32,
1873) -> u8 {
1874    // Validate scale_bits matches upstream constant
1875    // Panics on mismatch because this is a hot-path function that
1876    // should not return Result; upstream hardcodes this value.
1877    debug_assert_eq!(
1878        scale_bits, RANS_WORD_SCALE_BITS,
1879        "word rANS requires scale_bits={}",
1880        RANS_WORD_SCALE_BITS
1881    );
1882    let x = state.0;
1883    let mask = (1u32 << scale_bits) - 1;
1884    let slot = (x & mask) as usize;
1885    state.0 =
1886        (tables.slots[slot].freq as u32) * (x >> scale_bits) + (tables.slots[slot].bias as u32);
1887    tables.slot2sym[slot]
1888}
1889
1890/// Word rANS decoder renormalization.
1891///
1892/// Equivalent to `RansWordDecRenorm` in `rans_word_sse41.h`.
1893/// Reads a u16 word from the stream when the state falls below L.
1894/// At most 1 iteration because L=2^16 and state < L implies
1895/// (state << 16) | word < 2^32, which covers the full u32 range.
1896#[inline]
1897pub fn rans_word_dec_renorm(
1898    state: &mut RansWordState,
1899    reader: &mut Word16Reader,
1900) -> Result<(), DecodeError> {
1901    let x = state.0;
1902    if x < RANS_WORD_L {
1903        let w = reader.read_word16().ok_or(DecodeError::InputTooShort)? as u32;
1904        state.0 = (x << 16) | w;
1905    }
1906    Ok(())
1907}
1908
1909// ---------------------------------------------------------------------------
1910// Alias method — byte rANS with alias table
1911// ---------------------------------------------------------------------------
1912//
1913// The alias method converts a non-uniform frequency distribution into 256
1914// uniform buckets (each of size total/256). Each bucket may contain up to
1915// two symbols: a primary and an alias, separated by a divider threshold.
1916// This enables O(1) symbol lookup during decoding instead of binary search.
1917//
1918// Upstream source: main_alias.cpp (Vose's alias method + rANS)
1919//
1920// Requires the `alloc` feature for the alias_remap table.
1921
1922#[cfg(any(feature = "alloc", test))]
1923extern crate alloc as alloc_crate;
1924
1925/// Log2 of the number of symbols for alias method. Always 8 for byte rANS.
1926pub const ALIAS_LOG2_NSYMS: u32 = 8;
1927
1928/// Number of symbols for alias method. Always 256.
1929pub const ALIAS_NSYMS: usize = 1 << 8;
1930
1931/// Alias table for O(1) symbol decoding.
1932///
1933/// Constructed via Vose's algorithm from a normalized frequency distribution.
1934/// Each of the 256 buckets contains a threshold, with up to two symbols
1935/// (primary and alias) filling the bucket.
1936#[derive(Debug, Clone)]
1937#[cfg(any(feature = "alloc", test))]
1938pub struct AliasTable {
1939    /// Threshold per bucket: if xm < divider[i], use odd slot;
1940    /// otherwise use even slot. After construction, this stores
1941    /// the absolute position: `i * tgt_sum + original_divider`.
1942    pub divider: [u32; ALIAS_NSYMS],
1943    /// Frequencies for each of the 512 slots (2 per bucket).
1944    pub slot_freqs: [u32; ALIAS_NSYMS * 2],
1945    /// Adjustment values for state recovery during decode.
1946    pub slot_adjust: [u32; ALIAS_NSYMS * 2],
1947    /// Symbol IDs for each of the 512 slots.
1948    pub slot_symbols: [u8; ALIAS_NSYMS * 2],
1949    /// Remap table from cumulative-range slot to alias-table position.
1950    /// Sized to the total distribution (2^scale_bits entries, max 65536).
1951    pub alias_remap: alloc_crate::vec::Vec<u32>,
1952    /// The scale_bits used to construct this table.
1953    pub scale_bits: u32,
1954    /// Frequencies for each of the 256 symbols.
1955    pub freqs: [u32; ALIAS_NSYMS],
1956    /// Cumulative frequencies for each of the 256 symbols.
1957    pub cum_freqs: [u32; ALIAS_NSYMS + 1],
1958}
1959
1960/// Normalize frequencies to sum to exactly `target_total` (must be power of 2).
1961///
1962/// Equivalent to the normalization logic in `main_alias.cpp` `SymbolStats::normalize_freqs()`.
1963/// Resamples the distribution using cumulative frequency scaling and steals
1964/// frequency from other symbols if any non-zero symbol gets zeroed out.
1965#[inline]
1966pub fn rans_byte_alias_normalize_freqs(
1967    freqs: &[u32],
1968    num_symbols: usize,
1969    target_total: u32,
1970) -> Result<([u32; ALIAS_NSYMS], [u32; ALIAS_NSYMS + 1]), ModelError> {
1971    if num_symbols == 0 || num_symbols > ALIAS_NSYMS {
1972        return Err(ModelError::EmptyInput);
1973    }
1974    if target_total == 0 || (target_total & (target_total - 1)) != 0 {
1975        return Err(ModelError::InvalidScaleBits);
1976    }
1977
1978    let mut out_freqs = [0u32; ALIAS_NSYMS];
1979    let mut cum = [0u32; ALIAS_NSYMS + 1];
1980
1981    // Copy input frequencies, compute cumulative
1982    cum[0] = 0;
1983    for i in 0..num_symbols {
1984        out_freqs[i] = freqs[i];
1985        cum[i + 1] = cum[i]
1986            .checked_add(freqs[i])
1987            .ok_or(ModelError::TotalMismatch)?;
1988    }
1989    let cur_total = cum[num_symbols];
1990    if cur_total == 0 {
1991        return Err(ModelError::ZeroTotal);
1992    }
1993
1994    // Resample distribution based on cumulative freqs
1995    let target = target_total as u64;
1996    let cur_total_u64 = cur_total as u64;
1997    for i in 1..=num_symbols {
1998        cum[i] = ((target * cum[i] as u64) / cur_total_u64) as u32;
1999    }
2000
2001    // Zero-frequency theft: any non-zero symbol that was nuked to zero gets
2002    // frequency stolen from the symbol with the smallest frequency > 1.
2003    for i in 0..num_symbols {
2004        if out_freqs[i] != 0 && cum[i + 1] == cum[i] {
2005            // Find best symbol to steal from (smallest frequency > 1)
2006            let mut best_freq = u32::MAX;
2007            let mut best_steal = None;
2008            for j in 0..num_symbols {
2009                let freq = cum[j + 1] - cum[j];
2010                if freq > 1 && freq < best_freq {
2011                    best_freq = freq;
2012                    best_steal = Some(j);
2013                }
2014            }
2015            let steal = best_steal.ok_or(ModelError::WorkspaceTooSmall)?;
2016
2017            if steal < i {
2018                for j in (steal + 1)..=i {
2019                    cum[j] = cum[j].wrapping_sub(1);
2020                }
2021            } else {
2022                for j in (i + 1)..=steal {
2023                    cum[j] = cum[j].wrapping_add(1);
2024                }
2025            }
2026        }
2027    }
2028
2029    // Recompute frequencies from cumulative
2030    for i in 0..num_symbols {
2031        if out_freqs[i] == 0 {
2032            debug_assert!(cum[i + 1] == cum[i]);
2033        } else {
2034            debug_assert!(cum[i + 1] > cum[i]);
2035        }
2036        out_freqs[i] = cum[i + 1] - cum[i];
2037    }
2038
2039    Ok((out_freqs, cum))
2040}
2041
2042/// Build an alias table from a frequency distribution using Vose's algorithm.
2043///
2044/// Equivalent to `SymbolStats::make_alias_table()` in `main_alias.cpp`.
2045/// The frequencies must sum to exactly `1 << scale_bits`.
2046#[inline]
2047#[cfg(any(feature = "alloc", test))]
2048pub fn rans_byte_alias_build_table(
2049    freqs: &[u32; ALIAS_NSYMS],
2050    cum_freqs: &[u32; ALIAS_NSYMS + 1],
2051    scale_bits: u32,
2052) -> AliasTable {
2053    let total = 1u64 << scale_bits;
2054    let tgt_sum = (total / ALIAS_NSYMS as u64) as u32;
2055
2056    let mut divider = [0u32; ALIAS_NSYMS];
2057    let mut slot_freqs = [0u32; ALIAS_NSYMS * 2];
2058    let mut slot_adjust = [0u32; ALIAS_NSYMS * 2];
2059    let mut slot_symbols = [0u8; ALIAS_NSYMS * 2];
2060
2061    // ---- Phase 1: Vose's alias construction ----
2062
2063    let mut remaining = [0u32; ALIAS_NSYMS];
2064    for i in 0..ALIAS_NSYMS {
2065        remaining[i] = freqs[i];
2066        divider[i] = tgt_sum;
2067        slot_symbols[i * 2] = i as u8;
2068        slot_symbols[i * 2 + 1] = i as u8;
2069    }
2070
2071    // Find initial small and large buckets
2072    let mut cur_large = 0;
2073    while cur_large < ALIAS_NSYMS && remaining[cur_large] < tgt_sum {
2074        cur_large += 1;
2075    }
2076    let mut cur_small = 0;
2077    while cur_small < ALIAS_NSYMS && remaining[cur_small] >= tgt_sum {
2078        cur_small += 1;
2079    }
2080    let mut next_small = cur_small + 1;
2081
2082    // Top up small buckets from large buckets
2083    while cur_large < ALIAS_NSYMS && cur_small < ALIAS_NSYMS {
2084        slot_symbols[cur_small * 2] = cur_large as u8;
2085        divider[cur_small] = remaining[cur_small];
2086        remaining[cur_large] -= tgt_sum - divider[cur_small];
2087
2088        if remaining[cur_large] >= tgt_sum || next_small <= cur_large {
2089            cur_small = next_small;
2090            while cur_small < ALIAS_NSYMS && remaining[cur_small] >= tgt_sum {
2091                cur_small += 1;
2092            }
2093            next_small = cur_small + 1;
2094        } else {
2095            cur_small = cur_large;
2096        }
2097
2098        while cur_large < ALIAS_NSYMS && remaining[cur_large] < tgt_sum {
2099            cur_large += 1;
2100        }
2101    }
2102
2103    // ---- Phase 2: Distribute code slots ----
2104
2105    let mut assigned = [0u32; ALIAS_NSYMS];
2106    let total_slots = total as usize;
2107    let mut alias_remap = alloc_crate::vec![0u32; total_slots];
2108
2109    for i in 0..ALIAS_NSYMS {
2110        let j = slot_symbols[i * 2] as usize;
2111        let sym0_height = divider[i]; // the small symbol's amount
2112        let sym1_height = tgt_sum - divider[i]; // the large symbol's amount
2113        let base0 = assigned[i];
2114        let base1 = assigned[j];
2115        let cbase0 = cum_freqs[i] + base0;
2116        let cbase1 = cum_freqs[j] + base1;
2117
2118        divider[i] = i as u32 * tgt_sum + sym0_height;
2119
2120        slot_freqs[i * 2 + 1] = freqs[i];
2121        slot_freqs[i * 2] = freqs[j];
2122        slot_adjust[i * 2 + 1] = (i as u32 * tgt_sum).wrapping_sub(base0);
2123        slot_adjust[i * 2] = (i as u32 * tgt_sum).wrapping_sub(base1.wrapping_sub(sym0_height));
2124
2125        for k in 0..sym0_height {
2126            let idx = (cbase0 + k) as usize;
2127            if idx < total_slots {
2128                alias_remap[idx] = k + i as u32 * tgt_sum;
2129            }
2130        }
2131        for k in 0..sym1_height {
2132            let idx = (cbase1 + k) as usize;
2133            if idx < total_slots {
2134                alias_remap[idx] = (k + sym0_height) + i as u32 * tgt_sum;
2135            }
2136        }
2137
2138        assigned[i] += sym0_height;
2139        assigned[j] += sym1_height;
2140    }
2141
2142    AliasTable {
2143        divider,
2144        slot_freqs,
2145        slot_adjust,
2146        slot_symbols,
2147        alias_remap,
2148        scale_bits,
2149        freqs: *freqs,
2150        cum_freqs: *cum_freqs,
2151    }
2152}
2153
2154/// Alias method encoder put-symbol.
2155///
2156/// Equivalent to `RansEncPutAlias` in `main_alias.cpp`.
2157/// Renormalizes the state, then encodes symbol `s` using the alias remap table.
2158/// The encoded symbol uses division-based encoding: `((x/freq) << scale_bits) + alias_remap[...]`.
2159#[inline]
2160#[cfg(feature = "alloc")]
2161pub fn rans_byte_alias_enc_put<W: BackwardWriter>(
2162    state: &mut RansByteState,
2163    writer: &mut W,
2164    table: &AliasTable,
2165    s: u8,
2166    scale_bits: u32,
2167) -> Result<(), EncodeError> {
2168    let freq = table.freqs[s as usize];
2169    let x = rans_byte_enc_renorm(state.0, writer, freq, scale_bits)?;
2170    let slot = table.alias_remap[(x % freq + table.cum_freqs[s as usize]) as usize];
2171    state.0 = ((x / freq) << scale_bits) + slot;
2172    Ok(())
2173}
2174
2175/// Alias method decoder get-symbol.
2176///
2177/// Equivalent to `RansDecGetAlias` in `main_alias.cpp`.
2178/// Extracts the symbol from the state using the alias table.
2179/// Returns the symbol and the updated state.
2180#[inline]
2181#[cfg(feature = "alloc")]
2182pub fn rans_byte_alias_dec_get(
2183    state: RansByteState,
2184    table: &AliasTable,
2185    scale_bits: u32,
2186) -> (u8, RansByteState) {
2187    let x = state.0;
2188    let mask = (1u32 << scale_bits).wrapping_sub(1);
2189    let xm = x & mask;
2190    let bucket_id = (xm >> (scale_bits - ALIAS_LOG2_NSYMS)) as usize;
2191    let mut bucket2 = bucket_id * 2;
2192    if xm < table.divider[bucket_id] {
2193        bucket2 += 1;
2194    }
2195    let s = table.slot_symbols[bucket2];
2196    let new_x = table.slot_freqs[bucket2] * (x >> scale_bits) + xm - table.slot_adjust[bucket2];
2197    (s, RansByteState(new_x))
2198}
2199
2200/// Alias method decoder renormalization.
2201///
2202/// Delegates to the standard byte rANS renormalization (reads bytes).
2203/// Equivalent to `RansDecRenorm` in `rans_byte.h`.
2204#[inline]
2205#[cfg(feature = "alloc")]
2206pub fn rans_byte_alias_dec_renorm<R: ForwardReader>(
2207    state: &mut RansByteState,
2208    reader: &mut R,
2209) -> Result<(), DecodeError> {
2210    rans_byte_dec_renorm(state, reader)
2211}
2212
2213/// Alias method decoder advance (get symbol + renormalize).
2214/// Combines `RansDecGetAlias` + `RansDecRenorm` from the upstream.
2215#[inline]
2216#[cfg(feature = "alloc")]
2217pub fn rans_byte_alias_dec_advance<R: ForwardReader>(
2218    state: &mut RansByteState,
2219    reader: &mut R,
2220    table: &AliasTable,
2221    scale_bits: u32,
2222) -> Result<u8, DecodeError> {
2223    let (s, new_state) = rans_byte_alias_dec_get(*state, table, scale_bits);
2224    *state = new_state;
2225    rans_byte_alias_dec_renorm(state, reader)?;
2226    Ok(s)
2227}
2228
2229// ---------------------------------------------------------------------------
2230// Tests
2231// ---------------------------------------------------------------------------
2232
2233#[cfg(test)]
2234mod tests {
2235    use super::*;
2236    extern crate alloc;
2237
2238    #[test]
2239    fn test_state_init() {
2240        let s = RansByteState::new();
2241        assert_eq!(s.get(), RANS_BYTE_L);
2242    }
2243
2244    #[test]
2245    fn test_backward_writer_basic() {
2246        let mut buf = [0u8; 10];
2247        let pos;
2248        {
2249            let mut w = BackwardByteWriter::new(&mut buf);
2250            assert!(w.write_byte(0xAB).is_ok());
2251            assert_eq!(w.position(), 9);
2252            assert!(w.write_byte(0xCD).is_ok());
2253            pos = w.position();
2254        }
2255        assert_eq!(pos, 8);
2256        assert_eq!(buf[8], 0xCD);
2257        assert_eq!(buf[9], 0xAB);
2258        assert_eq!(buf[8..10], [0xCD, 0xAB]);
2259    }
2260
2261    #[test]
2262    fn test_backward_writer_full() {
2263        let mut buf = [0u8; 2];
2264        let mut w = BackwardByteWriter::new(&mut buf);
2265        assert!(w.write_byte(1).is_ok());
2266        assert!(w.write_byte(2).is_ok());
2267        assert!(w.write_byte(3).is_err());
2268    }
2269
2270    #[test]
2271    fn test_backward_writer_u32_le() {
2272        let mut buf = [0u8; 8];
2273        let pos1;
2274        {
2275            let mut w = BackwardByteWriter::new(&mut buf);
2276            assert!(w.write_u32_le(0x01020304).is_ok());
2277            pos1 = w.position();
2278            assert!(w.write_u32_le(0x05060708).is_ok());
2279        }
2280        assert_eq!(pos1, 4);
2281        assert_eq!(buf[4..8], [0x04, 0x03, 0x02, 0x01]);
2282        assert_eq!(buf[0..4], [0x08, 0x07, 0x06, 0x05]);
2283    }
2284
2285    #[test]
2286    fn test_forward_reader_basic() {
2287        let buf = [0x10, 0x20, 0x30, 0x40];
2288        let mut r = ByteReader::new(&buf);
2289        assert_eq!(r.read_byte(), Some(0x10));
2290        assert_eq!(r.read_byte(), Some(0x20));
2291        assert_eq!(r.position(), 2);
2292    }
2293
2294    #[test]
2295    fn test_forward_reader_u32_le() {
2296        let buf = [0x04, 0x03, 0x02, 0x01, 0x08, 0x07, 0x06, 0x05];
2297        let mut r = ByteReader::new(&buf);
2298        assert_eq!(r.read_u32_le(), Some(0x01020304));
2299        assert_eq!(r.read_u32_le(), Some(0x05060708));
2300        assert_eq!(r.read_u32_le(), None);
2301    }
2302
2303    #[test]
2304    fn test_enc_symbol_init() {
2305        // Simple case: freq=2, scale_bits=14
2306        let sym = RansByteEncSymbol::new(100, 2, 14).unwrap();
2307        assert!(sym.x_max > 0);
2308        assert!(sym.rcp_freq > 0);
2309        assert_eq!(sym.bias, 100);
2310        assert_eq!(sym.cmpl_freq, ((1u32 << 14) - 2) as u16);
2311        // For freq=2, shift=1, rcp_shift = shift-1 = 0. This is expected.
2312
2313        assert_eq!(sym.rcp_shift, 0);
2314    }
2315
2316    #[test]
2317    fn test_enc_symbol_init_freq_one() {
2318        // Special case: freq=1
2319        let sym = RansByteEncSymbol::new(100, 1, 14).unwrap();
2320        assert!(sym.x_max > 0);
2321        assert_eq!(sym.rcp_freq, !0u32);
2322        assert_eq!(sym.rcp_shift, 0);
2323        assert_eq!(sym.bias, 100 + (1u32 << 14) - 1);
2324    }
2325
2326    #[test]
2327    fn test_enc_symbol_init_max_freq() {
2328        // freq = (1 << scale_bits) - start
2329        let scale_bits = 14;
2330        let start = 0;
2331        let freq = 1u32 << scale_bits;
2332        let sym = RansByteEncSymbol::new(start, freq, scale_bits).unwrap();
2333        assert!(sym.x_max > 0);
2334    }
2335
2336    #[test]
2337    fn test_slice_backward_writer() {
2338        let mut buf = [0u8; 10];
2339        let mut writer = SliceBackwardWriter(&mut buf[..]);
2340        assert!(writer.write_byte(0xAB).is_ok());
2341        assert_eq!(writer.0.len(), 9);
2342        assert!(writer.write_byte(0xCD).is_ok());
2343        assert_eq!(writer.0.len(), 8);
2344        assert_eq!(buf[8], 0xCD);
2345        assert_eq!(buf[9], 0xAB);
2346    }
2347
2348    #[test]
2349    fn test_slice_forward_reader() {
2350        let buf = [0x10, 0x20, 0x30];
2351        let mut r = &buf[..];
2352        assert_eq!(r.read_byte(), Some(0x10));
2353        assert_eq!(r.read_byte(), Some(0x20));
2354        assert_eq!(r.read_byte(), Some(0x30));
2355        assert_eq!(r.read_byte(), None);
2356    }
2357
2358    #[test]
2359    fn test_roundtrip_single_symbol() {
2360        let scale_bits = 14;
2361        let symbols = [42u8; 100];
2362
2363        // Encode
2364        let mut out = [0u8; 1024];
2365        let mut writer = BackwardByteWriter::new(&mut out);
2366
2367        let esym = RansByteEncSymbol::new(0, 1u32 << scale_bits, scale_bits).unwrap();
2368        let mut state = RansByteState::new();
2369
2370        for _i in (0..symbols.len()).rev() {
2371            rans_byte_enc_put_symbol(&mut state, &mut writer, &esym).unwrap();
2372        }
2373        rans_byte_enc_flush(&state, &mut writer).unwrap();
2374
2375        let encoded = writer.encoded();
2376        assert!(!encoded.is_empty(), "encoded output should not be empty");
2377
2378        // Decode
2379        let mut reader = ByteReader::new(encoded);
2380        let dsym = RansByteDecSymbol::new(0, 1u32 << scale_bits).unwrap();
2381        let mut dec_state = rans_byte_dec_init(&mut reader).unwrap();
2382        // With one symbol occupying the entire [0, 1<<scale_bits) range,
2383        // all cum2sym slots map to that symbol (42).
2384        let cum2sym = [42u8; 1 << 14];
2385
2386        let mut output = alloc::vec![0u8; symbols.len()];
2387        for i in 0..symbols.len() {
2388            let cf = rans_byte_dec_get(&dec_state, scale_bits);
2389            let s = cum2sym[cf as usize];
2390            output[i] = s;
2391            rans_byte_dec_advance_symbol(&mut dec_state, &mut reader, &dsym, scale_bits).unwrap();
2392        }
2393
2394        assert_eq!(
2395            output,
2396            &symbols[..],
2397            "single-symbol round-trip should match"
2398        );
2399        assert_eq!(
2400            output,
2401            &symbols[..],
2402            "single-symbol round-trip should match"
2403        );
2404    }
2405
2406    #[test]
2407    fn test_roundtrip_two_symbols() {
2408        // Two symbols with frequencies 1 and 3 (scale_bits=2, total=4)
2409        let scale_bits = 2;
2410        let _total = 1u32 << scale_bits; // 4
2411        let freq0 = 1u32;
2412        let freq1 = 3u32;
2413
2414        let symbols: alloc::vec::Vec<u8> = (0..10).map(|i| (i % 2) as u8).collect();
2415
2416        // Encode using division-based path for correctness
2417        let mut out = [0u8; 1024];
2418        let mut writer = BackwardByteWriter::new(&mut out);
2419
2420        let mut state = RansByteState::new();
2421        for idx in (0..symbols.len()).rev() {
2422            let s = symbols[idx];
2423            let start = if s == 0 { 0 } else { freq0 };
2424            let freq = if s == 0 { freq0 } else { freq1 };
2425            rans_byte_enc_put(&mut state, &mut writer, start, freq, scale_bits).unwrap();
2426        }
2427        rans_byte_enc_flush(&state, &mut writer).unwrap();
2428
2429        let encoded = writer.encoded();
2430        assert!(!encoded.is_empty());
2431
2432        // Decode
2433        let mut reader = ByteReader::new(encoded);
2434        let dsym0 = RansByteDecSymbol::new(0, freq0).unwrap();
2435        let dsym1 = RansByteDecSymbol::new(freq0, freq1).unwrap();
2436
2437        let mut dec_state = rans_byte_dec_init(&mut reader).unwrap();
2438        let cum2sym = [0u8, 0u8, 1u8, 1u8]; // slots 0-1 -> sym0, slots 2-3 -> sym1
2439
2440        let mut output = alloc::vec![0u8; symbols.len()];
2441        for i in 0..symbols.len() {
2442            let cf = rans_byte_dec_get(&dec_state, scale_bits);
2443            let s = cum2sym[cf as usize] as usize;
2444            output[i] = s as u8;
2445            let dsym = if s == 0 { &dsym0 } else { &dsym1 };
2446            rans_byte_dec_advance_symbol(&mut dec_state, &mut reader, dsym, scale_bits).unwrap();
2447        }
2448
2449        assert_eq!(output, symbols, "two-symbol round-trip should match");
2450    }
2451
2452    #[test]
2453    fn test_slice_trait_roundtrip() {
2454        let scale_bits = 14;
2455        let symbols: alloc::vec::Vec<u8> = (0..50).map(|i| (i % 17) as u8).collect();
2456
2457        // Encode using slice writer (via SliceBackwardWriter)
2458        let mut out = [0u8; 1024];
2459        let mut writer = SliceBackwardWriter(&mut out[..]);
2460        let mut state = RansByteState::new();
2461
2462        // Build uniform freq model
2463        let total = 1u32 << scale_bits;
2464        let n_syms = 17u32;
2465        let base_freq = total / n_syms;
2466        for i in (0..symbols.len()).rev() {
2467            let s = symbols[i] as u32;
2468            let start = s * base_freq;
2469            let freq = base_freq;
2470            rans_byte_enc_put(&mut state, &mut writer, start, freq, scale_bits).unwrap();
2471        }
2472        rans_byte_enc_flush(&state, &mut writer).unwrap();
2473
2474        let used = writer.0.len();
2475        let encoded = &out[used..];
2476
2477        // Decode
2478        let mut reader = ByteReader::new(encoded);
2479        let mut dec_state = rans_byte_dec_init(&mut reader).unwrap();
2480
2481        let mut output = alloc::vec![0u8; symbols.len()];
2482        for i in 0..symbols.len() {
2483            let cf = rans_byte_dec_get(&dec_state, scale_bits);
2484            let s = cf / base_freq;
2485            output[i] = s as u8;
2486            let start = s * base_freq;
2487            rans_byte_dec_advance(&mut dec_state, &mut reader, start, base_freq, scale_bits)
2488                .unwrap();
2489        }
2490
2491        assert_eq!(output, symbols, "uniform-symbol round-trip should match");
2492    }
2493
2494    #[test]
2495    fn test_reciprocal_roundtrip() {
2496        let scale_bits = 14;
2497        let total = 1u32 << scale_bits;
2498        // Use frequencies that sum to total
2499        let freq0 = total / 3;
2500        let freq1 = total / 3;
2501        let freq2 = total - freq0 - freq1;
2502
2503        let esym0 = RansByteEncSymbol::new(0, freq0, scale_bits).unwrap();
2504        let esym1 = RansByteEncSymbol::new(freq0, freq1, scale_bits).unwrap();
2505        let esym2 = RansByteEncSymbol::new(freq0 + freq1, freq2, scale_bits).unwrap();
2506
2507        let dsym0 = RansByteDecSymbol::new(0, freq0).unwrap();
2508        let dsym1 = RansByteDecSymbol::new(freq0, freq1).unwrap();
2509        let dsym2 = RansByteDecSymbol::new(freq0 + freq1, freq2).unwrap();
2510
2511        let symbols: alloc::vec::Vec<u8> = (0..50).map(|i| (i % 3) as u8).collect();
2512
2513        // Encode with reciprocal fast path
2514        let mut out = [0u8; 1024];
2515        let mut writer = BackwardByteWriter::new(&mut out);
2516        let mut state = RansByteState::new();
2517
2518        for idx in (0..symbols.len()).rev() {
2519            let s = symbols[idx] as usize;
2520            let esym = match s {
2521                0 => &esym0,
2522                1 => &esym1,
2523                _ => &esym2,
2524            };
2525            rans_byte_enc_put_symbol(&mut state, &mut writer, esym).unwrap();
2526        }
2527        rans_byte_enc_flush(&state, &mut writer).unwrap();
2528        let encoded = writer.encoded();
2529
2530        // Decode with division-based path
2531        let mut reader = ByteReader::new(encoded);
2532        let mut dec_state = rans_byte_dec_init(&mut reader).unwrap();
2533
2534        let cum2sym: alloc::vec::Vec<u8> = (0..total as usize)
2535            .map(|i| {
2536                if i < freq0 as usize {
2537                    0
2538                } else if i < (freq0 + freq1) as usize {
2539                    1
2540                } else {
2541                    2
2542                }
2543            })
2544            .collect();
2545
2546        let mut output = alloc::vec![0u8; symbols.len()];
2547        for i in 0..symbols.len() {
2548            let cf = rans_byte_dec_get(&dec_state, scale_bits);
2549            let s = cum2sym[cf as usize] as usize;
2550            output[i] = s as u8;
2551            let dsym = match s {
2552                0 => &dsym0,
2553                1 => &dsym1,
2554                _ => &dsym2,
2555            };
2556            rans_byte_dec_advance_symbol(&mut dec_state, &mut reader, dsym, scale_bits).unwrap();
2557        }
2558
2559        assert_eq!(output, symbols, "reciprocal round-trip should match");
2560    }
2561
2562    #[test]
2563    fn test_interleaved_roundtrip() {
2564        let scale_bits = 14;
2565        let symbols: alloc::vec::Vec<u8> = (0..77).map(|i| (i % 7) as u8).collect();
2566
2567        // Build frequency model (approximate uniform)
2568        let total = 1u32 << scale_bits;
2569        let base_freq = total / 7;
2570        let esyms: alloc::vec::Vec<RansByteEncSymbol> = (0..7)
2571            .map(|i| RansByteEncSymbol::new(i * base_freq, base_freq, scale_bits).unwrap())
2572            .collect();
2573        let dsyms: alloc::vec::Vec<RansByteDecSymbol> = (0..7)
2574            .map(|i| RansByteDecSymbol::new(i * base_freq, base_freq).unwrap())
2575            .collect();
2576
2577        // Encode interleaved
2578        let mut out = [0u8; 2048];
2579        let mut writer = BackwardByteWriter::new(&mut out);
2580
2581        let mut s0 = RansByteState::new();
2582        let mut s1 = RansByteState::new();
2583
2584        let n = symbols.len();
2585        if n & 1 != 0 {
2586            let s = symbols[n - 1] as usize;
2587            rans_byte_enc_put_symbol(&mut s0, &mut writer, &esyms[s]).unwrap();
2588        }
2589
2590        let mut i = n & !1;
2591        while i > 0 {
2592            let s1_idx = symbols[i - 1] as usize;
2593            let s0_idx = symbols[i - 2] as usize;
2594            rans_byte_enc_put_symbol(&mut s1, &mut writer, &esyms[s1_idx]).unwrap();
2595            rans_byte_enc_put_symbol(&mut s0, &mut writer, &esyms[s0_idx]).unwrap();
2596            i = i.wrapping_sub(2);
2597        }
2598
2599        rans_byte_enc_flush(&s1, &mut writer).unwrap();
2600        rans_byte_enc_flush(&s0, &mut writer).unwrap();
2601
2602        let encoded = writer.encoded();
2603
2604        // Decode interleaved
2605        let mut reader = ByteReader::new(encoded);
2606        let mut d0 = rans_byte_dec_init(&mut reader).unwrap();
2607        let mut d1 = rans_byte_dec_init(&mut reader).unwrap();
2608
2609        let cum2sym: alloc::vec::Vec<u8> = (0..total as usize)
2610            .map(|i| (i / base_freq as usize) as u8)
2611            .collect();
2612
2613        let mut output = alloc::vec![0u8; n];
2614        let even_n = n & !1;
2615
2616        let mut pos = 0;
2617        while pos < even_n {
2618            let cf0 = rans_byte_dec_get(&d0, scale_bits);
2619            let s0 = cum2sym[cf0 as usize] as usize;
2620            let cf1 = rans_byte_dec_get(&d1, scale_bits);
2621            let s1 = cum2sym[cf1 as usize] as usize;
2622
2623            output[pos] = s0 as u8;
2624            output[pos + 1] = s1 as u8;
2625
2626            rans_byte_dec_advance_symbol_step(&mut d0, &dsyms[s0], scale_bits);
2627            rans_byte_dec_advance_symbol_step(&mut d1, &dsyms[s1], scale_bits);
2628            rans_byte_dec_renorm(&mut d0, &mut reader).unwrap();
2629            rans_byte_dec_renorm(&mut d1, &mut reader).unwrap();
2630
2631            pos += 2;
2632        }
2633
2634        if n & 1 != 0 {
2635            let cf0 = rans_byte_dec_get(&d0, scale_bits);
2636            let s0 = cum2sym[cf0 as usize] as usize;
2637            output[n - 1] = s0 as u8;
2638            rans_byte_dec_advance_symbol(&mut d0, &mut reader, &dsyms[s0], scale_bits).unwrap();
2639        }
2640
2641        assert_eq!(output, symbols, "interleaved round-trip should match");
2642    }
2643
2644    #[test]
2645    fn test_reciprocal_equals_division() {
2646        // For a range of frequencies and states, verify the reciprocal fast
2647        // path produces the same result as the division-based reference.
2648        let scale_bits = 14;
2649        let total = 1u32 << scale_bits;
2650
2651        // Division-based reference (single encode step, no renormalization)
2652        fn div_put(x: u32, start: u32, freq: u32, scale_bits: u32) -> u32 {
2653            ((x / freq) << scale_bits) + (x % freq) + start
2654        }
2655
2656        let test_freqs = [1, 2, 3, 5, 7, 10, 100, 1000, total / 2, total - 1];
2657
2658        for &freq in &test_freqs {
2659            let start = 0;
2660            let esym = RansByteEncSymbol::new(start, freq, scale_bits).unwrap();
2661
2662            let test_states = [
2663                RANS_BYTE_L,
2664                RANS_BYTE_L + 1,
2665                RANS_BYTE_L * 2,
2666                RANS_BYTE_L * 4,
2667                RANS_BYTE_L * 8,
2668                (1u32 << 31) - 1,
2669            ];
2670
2671            for &test_state in &test_states {
2672                if test_state >= esym.x_max {
2673                    // Would need renormalization; skip for now
2674                    continue;
2675                }
2676
2677                let expected = div_put(test_state, start, freq, scale_bits);
2678
2679                let mut state_fast = RansByteState(test_state);
2680                let mut temp = [0u8; 8];
2681                let mut w = BackwardByteWriter::new(&mut temp);
2682                rans_byte_enc_put_symbol(&mut state_fast, &mut w, &esym).unwrap();
2683
2684                assert_eq!(
2685                    state_fast.0, expected,
2686                    "reciprocal mismatch for freq={}, start={}, state={}",
2687                    freq, start, test_state
2688                );
2689            }
2690        }
2691    }
2692
2693    #[test]
2694    fn test_oracle_reciprocal_parameters() {
2695        // Verify against compiled C oracle output
2696
2697        // freq=10, start=0, scale_bits=14
2698        let sym = RansByteEncSymbol::new(0, 10, 14).unwrap();
2699        assert_eq!(sym.x_max, 1310720);
2700        assert_eq!(sym.rcp_freq, 3435973837);
2701        assert_eq!(sym.bias, 0);
2702        assert_eq!(sym.cmpl_freq as u32, 16374);
2703        assert_eq!(sym.rcp_shift as u32, 3);
2704
2705        // freq=1 special case, start=100, scale_bits=14
2706        let sym = RansByteEncSymbol::new(100, 1, 14).unwrap();
2707        assert_eq!(sym.x_max, 131072);
2708        assert_eq!(sym.rcp_freq, 4294967295);
2709        assert_eq!(sym.bias, 16483);
2710        assert_eq!(sym.cmpl_freq as u32, 16383);
2711        assert_eq!(sym.rcp_shift as u32, 0);
2712
2713        // freq=16384 (full total), start=0, scale_bits=14
2714        let sym = RansByteEncSymbol::new(0, 16384, 14).unwrap();
2715        assert_eq!(sym.x_max, 2147483648);
2716        assert_eq!(sym.rcp_freq, 2147483648);
2717        assert_eq!(sym.bias, 0);
2718        assert_eq!(sym.cmpl_freq as u32, 0);
2719        assert_eq!(sym.rcp_shift as u32, 13);
2720
2721        // freq=2, start=0, scale_bits=14
2722        let sym = RansByteEncSymbol::new(0, 2, 14).unwrap();
2723        assert_eq!(sym.cmpl_freq as u32, 16382);
2724        assert_eq!(sym.rcp_shift as u32, 0);
2725        assert!(sym.rcp_freq == 2147483648 || sym.rcp_freq > 0);
2726    }
2727
2728    #[test]
2729    fn test_reciprocal_freq_one() {
2730        let scale_bits = 14;
2731        let freq = 1u32;
2732        let start = 100;
2733        let esym = RansByteEncSymbol::new(start, freq, scale_bits).unwrap();
2734
2735        // The freq=1 special case should give: x_new = x * M + start
2736        fn expected(x: u32, start: u32, scale_bits: u32) -> u32 {
2737            x * (1u32 << scale_bits) + start
2738        }
2739
2740        let test_states = [RANS_BYTE_L, RANS_BYTE_L + 10, RANS_BYTE_L * 3, (1u32 << 30)];
2741
2742        for &test_state in &test_states {
2743            if test_state >= esym.x_max {
2744                continue;
2745            }
2746            let mut state = RansByteState(test_state);
2747            let mut tmp = [0u8; 8];
2748            let mut w = BackwardByteWriter::new(&mut tmp);
2749            rans_byte_enc_put_symbol(&mut state, &mut w, &esym).unwrap();
2750
2751            assert_eq!(
2752                state.0,
2753                expected(test_state, start, scale_bits),
2754                "freq=1 mismatch for state={}",
2755                test_state
2756            );
2757        }
2758    }
2759
2760    #[test]
2761    fn test_decoder_symbol_init() {
2762        let dsym = RansByteDecSymbol::new(100, 50).unwrap();
2763        assert_eq!(dsym.start, 100);
2764        assert_eq!(dsym.freq, 50);
2765    }
2766
2767    #[test]
2768    fn test_reader_exhaustion() {
2769        let buf = [1u8; 3];
2770        let mut reader = ByteReader::new(&buf);
2771        assert!(reader.read_byte().is_some());
2772        assert!(reader.read_byte().is_some());
2773        assert!(reader.read_byte().is_some());
2774        assert!(reader.read_byte().is_none());
2775        assert!(reader.read_u32_le().is_none());
2776    }
2777
2778    #[test]
2779    fn test_writer_exhaustion() {
2780        let mut buf = [0u8; 2];
2781        let mut writer = BackwardByteWriter::new(&mut buf);
2782        assert!(writer.write_u32_le(0x12345678).is_err());
2783        assert!(writer.write_byte(1).is_ok());
2784        assert!(writer.write_byte(2).is_ok());
2785        assert!(writer.write_byte(3).is_err());
2786    }
2787
2788    // -----------------------------------------------------------------------
2789    // 64-bit rANS tests
2790    // -----------------------------------------------------------------------
2791
2792    #[test]
2793    fn test_rans64_state_init() {
2794        let s = Rans64State::new();
2795        assert_eq!(s.get(), RANS64_L);
2796        assert_eq!(s, Rans64State::default());
2797    }
2798
2799    #[test]
2800    fn test_rans64_word32_writer_basic() {
2801        let mut buf = [0u8; 12];
2802        let pos;
2803        let words;
2804        {
2805            let mut w = BackwardWord32Writer::new(&mut buf);
2806            assert!(w.write_word32(0xDEADBEEF).is_ok());
2807            assert!(w.write_word32(0xCAFEBABE).is_ok());
2808            // Room for 1 more word (12 bytes = 3 words, wrote 2)
2809            assert!(w.write_word32(0x12345678).is_ok());
2810            pos = w.position();
2811            words = w.words_written();
2812        }
2813        assert_eq!(pos, 0);
2814        assert_eq!(words, 3);
2815        assert_eq!(buf[8..12], 0xDEADBEEFu32.to_le_bytes());
2816        assert_eq!(buf[4..8], 0xCAFEBABEu32.to_le_bytes());
2817        assert_eq!(buf[0..4], 0x12345678u32.to_le_bytes());
2818    }
2819
2820    #[test]
2821    fn test_rans64_word32_reader_basic() {
2822        let mut buf = [0u8; 12];
2823        let v0 = 0xDEADBEEFu32;
2824        let v1 = 0xCAFEBABEu32;
2825        let v2 = 0x12345678u32;
2826        buf[0..4].copy_from_slice(&v0.to_le_bytes());
2827        buf[4..8].copy_from_slice(&v1.to_le_bytes());
2828        buf[8..12].copy_from_slice(&v2.to_le_bytes());
2829
2830        let mut r = Word32Reader::new(&buf);
2831        assert_eq!(r.read_word32(), Some(v0));
2832        assert_eq!(r.read_word32(), Some(v1));
2833        assert_eq!(r.read_word32(), Some(v2));
2834        assert_eq!(r.read_word32(), None);
2835        assert_eq!(r.words_consumed(), 3);
2836    }
2837
2838    #[test]
2839    fn test_rans64_enc_symbol_init() {
2840        // Simple case: freq=2, scale_bits=14
2841        let sym = Rans64EncSymbol::new(100, 2, 14).unwrap();
2842        assert!(sym.x_max > 0);
2843        assert!(sym.rcp_freq > 0);
2844        assert_eq!(sym.bias, 100);
2845        assert_eq!(sym.cmpl_freq, ((1u32 << 14) - 2) as u32);
2846        assert_eq!(sym.rcp_shift, 0);
2847
2848        // Check x_max formula: ((RANS64_L >> scale_bits) << 32) * freq
2849        let expected_x_max = ((RANS64_L >> 14) << 32) * 2;
2850        assert_eq!(sym.x_max, expected_x_max);
2851    }
2852
2853    #[test]
2854    fn test_rans64_enc_symbol_init_freq_one() {
2855        let sym = Rans64EncSymbol::new(100, 1, 14).unwrap();
2856        assert!(sym.x_max > 0);
2857        assert_eq!(sym.rcp_freq, !0u64);
2858        assert_eq!(sym.rcp_shift, 0);
2859        assert_eq!(sym.bias, 100 + (1u64 << 14) - 1);
2860    }
2861
2862    #[test]
2863    fn test_rans64_enc_symbol_init_large_scale() {
2864        // scale_bits=31, large freq to exercise 128-bit division
2865        let scale_bits = 30;
2866        let start = 0;
2867        let freq = (1u32 << 29) + 1; // large, irregular freq
2868        let sym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
2869        assert!(sym.rcp_freq > 0);
2870        assert!(sym.x_max > 0);
2871        assert_eq!(sym.bias, 0);
2872    }
2873
2874    #[test]
2875    fn test_rans64_roundtrip_single_symbol_division() {
2876        // Encode and decode a single symbol using the division-based path
2877        let scale_bits = 14;
2878        let n = 50;
2879        let symbols = [99u8; 50];
2880
2881        // Encode
2882        let mut out = [0u8; 4096];
2883        let mut writer = BackwardWord32Writer::new(&mut out);
2884
2885        let mut state = Rans64State::new();
2886        for _i in (0..n).rev() {
2887            rans64_enc_put(
2888                &mut state,
2889                &mut writer,
2890                0,                  // start
2891                1u32 << scale_bits, // freq = total
2892                scale_bits,
2893            )
2894            .unwrap();
2895        }
2896        rans64_enc_flush(&state, &mut writer).unwrap();
2897
2898        let encoded = writer.encoded();
2899        assert!(
2900            encoded.len() >= 8,
2901            "encoded should have at least 8 bytes (2 words)"
2902        );
2903        assert!(!encoded.is_empty());
2904
2905        // Decode
2906        let mut reader = Word32Reader::new(encoded);
2907        let dsym = Rans64DecSymbol::new(0, 1u32 << scale_bits).unwrap();
2908        let mut dec_state = rans64_dec_init(&mut reader).unwrap();
2909        let cum2sym = [99u8; 1 << 14];
2910
2911        let mut output = alloc::vec![0u8; n];
2912        for i in 0..n {
2913            let cf = rans64_dec_get(&dec_state, scale_bits);
2914            let s = cum2sym[cf as usize];
2915            output[i] = s;
2916            rans64_dec_advance_symbol(&mut dec_state, &mut reader, &dsym, scale_bits).unwrap();
2917        }
2918
2919        assert_eq!(output, symbols, "64-bit single-symbol division round-trip");
2920    }
2921
2922    #[test]
2923    fn test_rans64_roundtrip_two_symbols_division() {
2924        // Two symbols with distinct frequencies, using 64-bit state
2925        let scale_bits = 14;
2926        let total = 1u32 << scale_bits;
2927        let freq0 = total / 4;
2928        let freq1 = total - freq0;
2929
2930        let symbols: alloc::vec::Vec<u8> = (0..30).map(|i| (i % 2) as u8).collect();
2931
2932        // Encode with division-based path
2933        let mut out = [0u8; 4096];
2934        let mut writer = BackwardWord32Writer::new(&mut out);
2935
2936        let mut state = Rans64State::new();
2937        for idx in (0..symbols.len()).rev() {
2938            let s = symbols[idx];
2939            let start = if s == 0 { 0 } else { freq0 };
2940            let freq = if s == 0 { freq0 } else { freq1 };
2941            rans64_enc_put(&mut state, &mut writer, start, freq, scale_bits).unwrap();
2942        }
2943        rans64_enc_flush(&state, &mut writer).unwrap();
2944
2945        let encoded = writer.encoded();
2946        assert!(encoded.len() >= 8, "encoded length = {}", encoded.len());
2947
2948        // Decode
2949        let mut reader = Word32Reader::new(encoded);
2950        let dsym0 = Rans64DecSymbol::new(0, freq0).unwrap();
2951        let dsym1 = Rans64DecSymbol::new(freq0, freq1).unwrap();
2952
2953        let mut dec_state = rans64_dec_init(&mut reader).unwrap();
2954        // Build cum2sym for these two symbols
2955        let mut cum2sym = alloc::vec![1u8; total as usize];
2956        for i in 0..freq0 as usize {
2957            cum2sym[i] = 0;
2958        }
2959
2960        let mut output = alloc::vec![0u8; symbols.len()];
2961        for i in 0..symbols.len() {
2962            let cf = rans64_dec_get(&dec_state, scale_bits);
2963            let s = cum2sym[cf as usize] as usize;
2964            output[i] = s as u8;
2965            let dsym = if s == 0 { &dsym0 } else { &dsym1 };
2966            rans64_dec_advance_symbol(&mut dec_state, &mut reader, dsym, scale_bits).unwrap();
2967        }
2968
2969        assert_eq!(output, symbols, "64-bit two-symbol division round-trip");
2970    }
2971
2972    #[test]
2973    fn test_rans64_reciprocal_equals_division() {
2974        // Verify reciprocal fast path produces the same C(s, x) as division
2975        let scale_bits = 14;
2976        let total = 1u32 << scale_bits;
2977
2978        fn div_put(x: u64, start: u32, freq: u32, scale_bits: u32) -> u64 {
2979            ((x / (freq as u64)) << scale_bits) + (x % (freq as u64)) + (start as u64)
2980        }
2981
2982        let test_freqs = [1, 2, 3, 5, 7, 10, 100, 1000, total / 2, total - 1];
2983
2984        for &freq in &test_freqs {
2985            let start = 0u32;
2986            let esym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
2987
2988            let test_states = [
2989                RANS64_L,
2990                RANS64_L + 1,
2991                RANS64_L * 2,
2992                RANS64_L * 4,
2993                RANS64_L * 8,
2994                (1u64 << 62) - 1,
2995            ];
2996
2997            for &test_state in &test_states {
2998                if test_state >= esym.x_max {
2999                    continue;
3000                }
3001
3002                let expected = div_put(test_state, start, freq, scale_bits);
3003
3004                let mut state_fast = Rans64State(test_state);
3005                let mut temp = [0u8; 16];
3006                let mut w = BackwardWord32Writer::new(&mut temp);
3007                rans64_enc_put_symbol(&mut state_fast, &mut w, &esym).unwrap();
3008
3009                assert_eq!(
3010                    state_fast.0, expected,
3011                    "64-bit reciprocal mismatch for freq={}, start={}, state={}",
3012                    freq, start, test_state
3013                );
3014            }
3015        }
3016    }
3017
3018    #[test]
3019    fn test_rans64_roundtrip_reciprocal() {
3020        // Full round-trip using reciprocal fast path for encoding,
3021        // division-based decoding
3022        let scale_bits = 14;
3023        let total = 1u32 << scale_bits;
3024        let freq0 = total / 3;
3025        let freq1 = total / 3;
3026        let freq2 = total - freq0 - freq1;
3027
3028        let esym0 = Rans64EncSymbol::new(0, freq0, scale_bits).unwrap();
3029        let esym1 = Rans64EncSymbol::new(freq0, freq1, scale_bits).unwrap();
3030        let esym2 = Rans64EncSymbol::new(freq0 + freq1, freq2, scale_bits).unwrap();
3031
3032        let dsym0 = Rans64DecSymbol::new(0, freq0).unwrap();
3033        let dsym1 = Rans64DecSymbol::new(freq0, freq1).unwrap();
3034        let dsym2 = Rans64DecSymbol::new(freq0 + freq1, freq2).unwrap();
3035
3036        let symbols: alloc::vec::Vec<u8> = (0..50).map(|i| (i % 3) as u8).collect();
3037
3038        // Encode with reciprocal fast path
3039        let mut out = [0u8; 4096];
3040        let mut writer = BackwardWord32Writer::new(&mut out);
3041        let mut state = Rans64State::new();
3042
3043        for idx in (0..symbols.len()).rev() {
3044            let s = symbols[idx] as usize;
3045            let esym = match s {
3046                0 => &esym0,
3047                1 => &esym1,
3048                _ => &esym2,
3049            };
3050            rans64_enc_put_symbol(&mut state, &mut writer, esym).unwrap();
3051        }
3052        rans64_enc_flush(&state, &mut writer).unwrap();
3053        let encoded = writer.encoded();
3054
3055        // Decode with division-based path
3056        let mut reader = Word32Reader::new(encoded);
3057        let mut dec_state = rans64_dec_init(&mut reader).unwrap();
3058
3059        let cum2sym: alloc::vec::Vec<u8> = (0..total as usize)
3060            .map(|i| {
3061                if i < freq0 as usize {
3062                    0
3063                } else if i < (freq0 + freq1) as usize {
3064                    1
3065                } else {
3066                    2
3067                }
3068            })
3069            .collect();
3070
3071        let mut output = alloc::vec![0u8; symbols.len()];
3072        for i in 0..symbols.len() {
3073            let cf = rans64_dec_get(&dec_state, scale_bits);
3074            let s = cum2sym[cf as usize] as usize;
3075            output[i] = s as u8;
3076            let dsym = match s {
3077                0 => &dsym0,
3078                1 => &dsym1,
3079                _ => &dsym2,
3080            };
3081            rans64_dec_advance_symbol(&mut dec_state, &mut reader, dsym, scale_bits).unwrap();
3082        }
3083
3084        assert_eq!(output, symbols, "64-bit reciprocal round-trip");
3085    }
3086
3087    #[test]
3088    fn test_rans64_step_operations() {
3089        // Verify the step-only operations produce the same intermediate state
3090        let scale_bits = 14;
3091        let total = 1u32 << scale_bits;
3092        let freq = total / 2;
3093        let start = 0;
3094
3095        let dsym = Rans64DecSymbol::new(start, freq).unwrap();
3096
3097        // Start with a state large enough to not need renormalization
3098        let state_val = RANS64_L * 4;
3099        let mut state_advance = Rans64State(state_val);
3100        let mut state_step = Rans64State(state_val);
3101
3102        // Advance via regular advance (no renorm needed since x >= L)
3103        // We need a dummy buffer with enough words to not actually use them
3104        let dummy_buf = [0u8; 16];
3105        let mut reader = Word32Reader::new(&dummy_buf);
3106        rans64_dec_advance(&mut state_advance, &mut reader, start, freq, scale_bits).unwrap();
3107
3108        // Advance via step-only
3109        rans64_dec_advance_step(&mut state_step, start, freq, scale_bits);
3110
3111        assert_eq!(
3112            state_advance.0, state_step.0,
3113            "step-only advance should match regular advance when no renorm needed"
3114        );
3115
3116        // Repeat with symbol convenience
3117        let mut state_adv_sym = Rans64State(state_val);
3118        let mut state_step_sym = Rans64State(state_val);
3119        let mut reader2 = Word32Reader::new(&dummy_buf);
3120        rans64_dec_advance_symbol(&mut state_adv_sym, &mut reader2, &dsym, scale_bits).unwrap();
3121        rans64_dec_advance_symbol_step(&mut state_step_sym, &dsym, scale_bits);
3122
3123        assert_eq!(
3124            state_adv_sym.0, state_step_sym.0,
3125            "step-only symbol advance should match regular"
3126        );
3127    }
3128
3129    #[test]
3130    fn test_rans64_state_transition_cycle() {
3131        // Encode a symbol and verify the decode retrieves the original
3132        // Uses a well-known state transition: C(s, x) then D(s, C(s, x))
3133        let scale_bits = 14;
3134        let _total = 1u32 << scale_bits;
3135        let freq = 100u32;
3136        let start = 500u32;
3137
3138        let esym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
3139        let _dsym = Rans64DecSymbol::new(start, freq).unwrap();
3140
3141        // Pick a test state that doesn't need renormalization
3142        let x = RANS64_L; // minimum valid state
3143
3144        // Encode step: C(s, x) using reciprocal
3145        let mut enc_state = Rans64State(x);
3146        let mut tmp = [0u8; 16];
3147        let mut w = BackwardWord32Writer::new(&mut tmp);
3148        rans64_enc_put_symbol(&mut enc_state, &mut w, &esym).unwrap();
3149        let encoded_x = enc_state.0;
3150
3151        // Decode step: D(s, C(s, x)) should give back x
3152        let dummy = [0u8; 16];
3153        let mut r = Word32Reader::new(&dummy);
3154        let mut dec_state = Rans64State(encoded_x);
3155        rans64_dec_advance(&mut dec_state, &mut r, start, freq, scale_bits).unwrap();
3156
3157        assert_eq!(
3158            dec_state.0, x,
3159            "decoding should invert encoding: D(s, C(s, x)) = x"
3160        );
3161    }
3162
3163    #[test]
3164    fn test_rans64_flush_init_roundtrip() {
3165        // Verify that flushing a state and re-initializing from those words
3166        // gives back the same state
3167        let test_state = 0xDEADBEEF_CAFEBABEu64;
3168        let state_in = Rans64State(test_state);
3169
3170        // Flush: write 2 u32 words (low first, then high)
3171        let mut buf = [0u8; 16];
3172        let mut writer = BackwardWord32Writer::new(&mut buf);
3173        rans64_enc_flush(&state_in, &mut writer).unwrap();
3174
3175        let encoded = writer.encoded();
3176        assert_eq!(encoded.len(), 8, "flush should write exactly 8 bytes");
3177
3178        // Verify the byte layout directly:
3179        // Low word is stored at position 0..4, high word at 4..8
3180        let lo_expected = (test_state & 0xffffffff) as u32;
3181        let hi_expected = (test_state >> 32) as u32;
3182        let lo_actual = u32::from_le_bytes([encoded[0], encoded[1], encoded[2], encoded[3]]);
3183        let hi_actual = u32::from_le_bytes([encoded[4], encoded[5], encoded[6], encoded[7]]);
3184        assert_eq!(lo_actual, lo_expected, "low word should match");
3185        assert_eq!(hi_actual, hi_expected, "high word should match");
3186
3187        // Re-init: read back the state
3188        let mut reader = Word32Reader::new(encoded);
3189        let state_out = rans64_dec_init(&mut reader).unwrap();
3190        assert_eq!(state_out.0, test_state, "flush+init round-trip");
3191    }
3192
3193    #[test]
3194    fn test_rans64_mul_hi() {
3195        // Verify rans64_mul_hi against u128 reference
3196        // 0xABCDEF0123456789 * 0x9876543210FEDCBA
3197        let a = 0xABCDEF0123456789u64;
3198        let b = 0x9876543210FEDCBAu64;
3199        let expected = ((a as u128) * (b as u128) >> 64) as u64;
3200        assert_eq!(rans64_mul_hi(a, b), expected);
3201
3202        // Simple cases
3203        assert_eq!(rans64_mul_hi(1, 1), 0);
3204        assert_eq!(rans64_mul_hi(1u64 << 63, 2), 1);
3205        assert_eq!(rans64_mul_hi(!0u64, !0u64), !0u64 - 1);
3206    }
3207
3208    #[test]
3209    fn test_rans64_decoder_symbol_init() {
3210        let dsym = Rans64DecSymbol::new(100, 50).unwrap();
3211        assert_eq!(dsym.start, 100);
3212        assert_eq!(dsym.freq, 50);
3213    }
3214
3215    #[test]
3216    fn test_rans64_freq_one_special() {
3217        // For freq=1, the reciprocal path should match: x * M + start
3218        let scale_bits = 14;
3219        let freq = 1u32;
3220        let start = 100;
3221        let esym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
3222
3223        fn expected(x: u64, start: u64, scale_bits: u32) -> u64 {
3224            x * (1u64 << scale_bits) + start
3225        }
3226
3227        let test_states = [RANS64_L, RANS64_L + 10, RANS64_L * 3, (1u64 << 60)];
3228
3229        for &test_state in &test_states {
3230            if test_state >= esym.x_max {
3231                continue;
3232            }
3233            let mut state = Rans64State(test_state);
3234            let mut tmp = [0u8; 16];
3235            let mut w = BackwardWord32Writer::new(&mut tmp);
3236            rans64_enc_put_symbol(&mut state, &mut w, &esym).unwrap();
3237
3238            assert_eq!(
3239                state.0,
3240                expected(test_state, start as u64, scale_bits),
3241                "64-bit freq=1 mismatch for state={}",
3242                test_state
3243            );
3244        }
3245    }
3246
3247    #[test]
3248    fn test_rans64_word32_writer_exhaustion() {
3249        let mut buf = [0u8; 4]; // only room for 1 word
3250        let mut writer = BackwardWord32Writer::new(&mut buf);
3251        assert!(writer.write_word32(0x12345678).is_ok());
3252        assert!(writer.write_word32(0x9ABCDEF0).is_err());
3253    }
3254
3255    #[test]
3256    fn test_rans64_word32_reader_exhaustion() {
3257        let buf = [0x01, 0x02, 0x03]; // only 3 bytes, not enough for one word
3258        let mut reader = Word32Reader::new(&buf);
3259        assert!(reader.read_word32().is_none());
3260    }
3261
3262    #[test]
3263    fn test_rans64_renorm_roundtrip() {
3264        // Encode a symbol large enough to trigger renormalization,
3265        // then decode it back using the full path
3266        let scale_bits = 14;
3267        let total = 1u32 << scale_bits;
3268
3269        // Use a small freq so x_max is small, forcing renormalization
3270        let freq = 7u32;
3271        let start = 100;
3272        let esym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
3273        let dsym = Rans64DecSymbol::new(start, freq).unwrap();
3274
3275        let mut out = [0u8; 4096];
3276        let mut writer = BackwardWord32Writer::new(&mut out);
3277        let mut state = Rans64State::new();
3278
3279        // Encode many symbols to force renormalization
3280        let n = 100;
3281        for _i in 0..n {
3282            rans64_enc_put_symbol(&mut state, &mut writer, &esym).unwrap();
3283        }
3284        rans64_enc_flush(&state, &mut writer).unwrap();
3285
3286        let encoded = writer.encoded();
3287
3288        // Decode
3289        let mut reader = Word32Reader::new(encoded);
3290        let mut dec_state = rans64_dec_init(&mut reader).unwrap();
3291
3292        // Build cum2sym mapping for this single symbol range
3293        let cum2sym: alloc::vec::Vec<u8> = (0..total as usize)
3294            .map(|i| {
3295                if (i as u32) < start {
3296                    255 // shouldn't happen
3297                } else if (i as u32) < start + freq {
3298                    42
3299                } else {
3300                    255 // shouldn't happen
3301                }
3302            })
3303            .collect();
3304
3305        let mut output = alloc::vec![0u8; n];
3306        for i in 0..n {
3307            let cf = rans64_dec_get(&dec_state, scale_bits);
3308            let s = cum2sym[cf as usize];
3309            output[i] = s;
3310            rans64_dec_advance_symbol(&mut dec_state, &mut reader, &dsym, scale_bits).unwrap();
3311        }
3312
3313        assert_eq!(output.len(), n);
3314        for &val in &output {
3315            assert_eq!(val, 42, "all decoded symbols should be 42");
3316        }
3317    }
3318
3319    #[test]
3320    fn test_rans64_renorm_only() {
3321        // Test rans64_dec_renorm in isolation with a prepared reader
3322        // Decoder state below RANS64_L, check that reading words brings it back up
3323        let mut buf = [0u8; 12];
3324        // Write two u32 words that, when shifted in, will push state >= RANS64_L
3325        let w0 = 0x00000001u32;
3326        let w1 = 0x00000002u32;
3327        buf[0..4].copy_from_slice(&w0.to_le_bytes());
3328        buf[4..8].copy_from_slice(&w1.to_le_bytes());
3329
3330        let mut reader = Word32Reader::new(&buf);
3331
3332        // State below L: reading one word should push it above L
3333        let mut state = Rans64State(RANS64_L - 1);
3334        rans64_dec_renorm(&mut state, &mut reader).unwrap();
3335        assert!(
3336            state.0 >= RANS64_L,
3337            "after renorm, state {} should be >= RANS64_L",
3338            state.0
3339        );
3340        // Specifically: (RANS64_L - 1) << 32 | 1 >= RANS64_L
3341        assert_eq!(state.0, ((RANS64_L - 1) << 32) | 1);
3342        assert_eq!(reader.words_consumed(), 1);
3343    }
3344
3345    #[test]
3346    fn test_rans64_large_scale_reciprocal() {
3347        // Test 64-bit reciprocal parameters across scale_bits 17..31
3348        // Verifying that cmpl_freq (u32) correctly handles values > 65535.
3349        use super::*;
3350
3351        for scale_bits in 17u32..=31u32 {
3352            let total = 1u64 << scale_bits;
3353
3354            // Pick frequencies that produce complement frequencies > u16::MAX
3355            let test_cases = [
3356                (0u32, 100u32),                       // small freq, start=0
3357                (0u32, 50000u32),                     // freq that needs >16-bit cmpl
3358                (100u32, 1u32),                       // freq=1 special case
3359                (0u32, (1u32 << scale_bits.min(20))), // large freq
3360                (total as u32 / 3, total as u32 / 3), // start > 0, freq > 0
3361            ];
3362
3363            for &(start, freq) in &test_cases {
3364                // Skip invalid cases
3365                if freq == 0 {
3366                    continue;
3367                }
3368                if (start as u64) + (freq as u64) > total {
3369                    continue;
3370                }
3371
3372                let sym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
3373
3374                // Verify cmpl_freq is correct (u32, not truncated to u16)
3375                let expected_cmpl = ((1u64 << scale_bits) - freq as u64) as u32;
3376                assert_eq!(
3377                    sym.cmpl_freq, expected_cmpl,
3378                    "cmpl_freq mismatch for scale_bits={}, start={}, freq={}: expected {}, got {}",
3379                    scale_bits, start, freq, expected_cmpl, sym.cmpl_freq
3380                );
3381
3382                // Verify x_max matches upstream formula
3383                let expected_x_max = ((RANS64_L >> scale_bits) << 32) * (freq as u64);
3384                assert_eq!(
3385                    sym.x_max, expected_x_max,
3386                    "x_max mismatch for scale_bits={}, start={}, freq={}",
3387                    scale_bits, start, freq
3388                );
3389
3390                // For freq >= 2, verify that rcp_freq * freq approximately equals 2^(shift+63)
3391                if freq >= 2 {
3392                    assert!(sym.rcp_freq > 0, "rcp_freq must be > 0 for freq={}", freq);
3393
3394                    // Verify rcp_shift consistent with freq
3395                    let mut expected_shift = 0u32;
3396                    while freq > (1u32 << expected_shift) {
3397                        expected_shift += 1;
3398                    }
3399                    assert_eq!(
3400                        sym.rcp_shift,
3401                        expected_shift - 1,
3402                        "rcp_shift mismatch for freq={}",
3403                        freq
3404                    );
3405                }
3406
3407                // Verify bias
3408                let expected_bias = if freq < 2 {
3409                    (start as u64) + (1u64 << scale_bits) - 1
3410                } else {
3411                    start as u64
3412                };
3413                assert_eq!(
3414                    sym.bias, expected_bias,
3415                    "bias mismatch for scale_bits={}, start={}, freq={}",
3416                    scale_bits, start, freq
3417                );
3418            }
3419        }
3420    }
3421
3422    #[test]
3423    fn test_rans64_reciprocal_equals_division_large() {
3424        // Verify that the reciprocal fast path produces the same state
3425        // as the division-based reference for large scale_bits and various states.
3426        use super::*;
3427
3428        let scale_bits = 30;
3429
3430        // Frequencies that produce >16-bit complements
3431        let freqs = [1u32, 2, 100, 10000, 500000000, 1000000000, (1u32 << 30) - 1];
3432        let total = 1u64 << scale_bits;
3433
3434        for &freq in &freqs {
3435            let start = 0u32;
3436            if (start as u64) + (freq as u64) > total {
3437                continue;
3438            }
3439
3440            let sym = Rans64EncSymbol::new(start, freq, scale_bits).unwrap();
3441
3442            // Test several state values that are within normalization bounds
3443            let states = [RANS64_L, RANS64_L + 1, RANS64_L * 2, (1u64 << 62) - 1];
3444
3445            for &state_val in &states {
3446                if state_val >= sym.x_max {
3447                    continue; // skip states that would trigger renormalization
3448                }
3449
3450                // Division-based reference
3451                let div_state = ((state_val / freq as u64) << scale_bits)
3452                    + (state_val % freq as u64)
3453                    + start as u64;
3454
3455                // Reciprocal fast path
3456                let q = rans64_mul_hi(state_val, sym.rcp_freq) >> sym.rcp_shift;
3457                let fast_state = state_val + sym.bias + q * (sym.cmpl_freq as u64);
3458
3459                assert_eq!(
3460                    fast_state, div_state,
3461                    "reciprocal mismatch for scale_bits={}, freq={}, state={}: div={}, fast={}",
3462                    scale_bits, freq, state_val, div_state, fast_state
3463                );
3464            }
3465        }
3466    }
3467}