Skip to main content

unarc_rs/ha/
acoder.rs

1//! Arithmetic Decoder for HA archive decompression
2//!
3//! Implements a 16-bit range coder with E3 (underflow) scaling.
4//! Based on the standard arithmetic coding algorithm as described in:
5//! - Witten, Neal, Cleary: "Arithmetic Coding for Data Compression" (1987)
6//! - Moffat, Neal, Witten: "Arithmetic Coding Revisited" (1998)
7
8use std::io::{BufReader, Read};
9
10use crate::error::Result;
11
12/// Full 16-bit range for interval arithmetic
13const RANGE_MAX: u16 = 0xFFFF;
14
15/// MSB mask for convergence check
16const MSB_MASK: u16 = 0x8000;
17
18/// Second-MSB mask for underflow (E3) check
19const UNDERFLOW_MASK: u16 = 0x4000;
20
21/// Arithmetic decoder using 16-bit precision with E3 underflow prevention.
22///
23/// The decoder maintains an interval [low, high] and a code value read from
24/// the bitstream. As symbols are decoded, the interval narrows and bits are
25/// shifted out when the MSBs converge.
26pub struct ArithmeticDecoder<R: Read> {
27    /// Buffered input stream
28    input: BufReader<R>,
29    /// Upper bound of current interval
30    high: u16,
31    /// Lower bound of current interval
32    low: u16,
33    /// Current code value from bitstream
34    code: u16,
35    /// Buffered byte for bit extraction
36    byte_buffer: u8,
37    /// Remaining bits in buffer (0-8)
38    bits_remaining: u8,
39}
40
41impl<R: Read> ArithmeticDecoder<R> {
42    /// Initialize decoder by reading the first 16 bits of encoded data.
43    #[inline]
44    pub fn new(reader: R) -> Result<Self> {
45        let mut input = BufReader::new(reader);
46
47        // Bootstrap: read initial 16-bit code value (big-endian)
48        let mut initial_bytes = [0u8; 2];
49        input.read_exact(&mut initial_bytes)?;
50        let code = u16::from_be_bytes(initial_bytes);
51
52        Ok(Self {
53            input,
54            high: RANGE_MAX,
55            low: 0,
56            code,
57            byte_buffer: 0,
58            bits_remaining: 0,
59        })
60    }
61
62    /// Extract next bit from input stream (MSB first within each byte).
63    #[inline(always)]
64    fn read_bit(&mut self) -> u16 {
65        if self.bits_remaining == 0 {
66            let mut byte = [0u8; 1];
67            // On EOF, feed zeros (standard practice for trailing bits)
68            if self.input.read_exact(&mut byte).is_ok() {
69                self.byte_buffer = byte[0];
70            } else {
71                self.byte_buffer = 0;
72            }
73            self.bits_remaining = 8;
74        }
75
76        self.bits_remaining -= 1;
77        ((self.byte_buffer >> self.bits_remaining) & 1) as u16
78    }
79
80    /// Compute threshold for symbol lookup given total frequency.
81    ///
82    /// Returns a value in [0, total-1] indicating where the current code
83    /// falls within the frequency distribution.
84    #[inline]
85    pub fn threshold_val(&self, total: u16) -> u16 {
86        // Interval width = high - low + 1
87        let range = (self.high - self.low) as u32 + 1;
88        // Position of code within interval, scaled to [0, total)
89        let offset = (self.code - self.low) as u32 + 1;
90        ((offset * total as u32 - 1) / range) as u16
91    }
92
93    /// Narrow interval after decoding a symbol and renormalize.
94    ///
95    /// - `cum_low`: cumulative frequency of all symbols before this one
96    /// - `cum_high`: cumulative frequency up to and including this symbol
97    /// - `total`: sum of all symbol frequencies
98    #[inline]
99    pub fn decode_update(&mut self, cum_low: u16, cum_high: u16, total: u16) -> Result<()> {
100        let range = (self.high - self.low) as u32 + 1;
101        let scale = total as u32;
102
103        // Narrow interval to [new_low, new_high]
104        let new_high = self.low.wrapping_add(((range * cum_high as u32 / scale) - 1) as u16);
105        let new_low = self.low.wrapping_add((range * cum_low as u32 / scale) as u16);
106
107        self.high = new_high;
108        self.low = new_low;
109
110        // Renormalization loop
111        self.renormalize();
112
113        Ok(())
114    }
115
116    /// Shift out converged bits and handle E3 underflow scaling.
117    #[inline(always)]
118    fn renormalize(&mut self) {
119        loop {
120            if (self.high ^ self.low) & MSB_MASK == 0 {
121                // MSBs match: shift out and read new bit
122                self.shift_out_msb();
123            } else if (self.low & UNDERFLOW_MASK) != 0 && (self.high & UNDERFLOW_MASK) == 0 {
124                // E3 underflow: low = 01..., high = 10...
125                // Expand interval around midpoint
126                self.handle_underflow();
127            } else {
128                // Interval is normalized
129                break;
130            }
131        }
132    }
133
134    /// Shift out matching MSB and read next bit into code.
135    #[inline(always)]
136    fn shift_out_msb(&mut self) {
137        self.low <<= 1;
138        self.high = (self.high << 1) | 1;
139        self.code = (self.code << 1) | self.read_bit();
140    }
141
142    /// Handle E3 scaling by flipping MSB and shifting.
143    #[inline(always)]
144    fn handle_underflow(&mut self) {
145        // Clear bit 14 of low, set bit 15 and bit 0 of high
146        self.low = (self.low << 1) & 0x7FFF;
147        self.high = (self.high << 1) | 0x8001;
148        // Flip MSB of code and shift in new bit
149        self.code = ((self.code << 1) ^ MSB_MASK) | self.read_bit();
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use std::io::Cursor;
157
158    #[test]
159    fn test_initialization() {
160        let data = vec![0xAB, 0xCD];
161        let decoder = ArithmeticDecoder::new(Cursor::new(data)).unwrap();
162
163        assert_eq!(decoder.low, 0);
164        assert_eq!(decoder.high, 0xFFFF);
165        assert_eq!(decoder.code, 0xABCD);
166    }
167
168    #[test]
169    fn test_threshold_midpoint() {
170        // Code at exact midpoint of full range
171        let data = vec![0x80, 0x00];
172        let decoder = ArithmeticDecoder::new(Cursor::new(data)).unwrap();
173
174        // threshold = ((0x8000 + 1) * 256 - 1) / 0x10000 = 128
175        assert_eq!(decoder.threshold_val(256), 128);
176    }
177
178    #[test]
179    fn test_threshold_boundaries() {
180        // Code at minimum
181        let data = vec![0x00, 0x00];
182        let decoder = ArithmeticDecoder::new(Cursor::new(data)).unwrap();
183        assert_eq!(decoder.threshold_val(100), 0);
184
185        // Code at maximum
186        let data = vec![0xFF, 0xFF];
187        let decoder = ArithmeticDecoder::new(Cursor::new(data)).unwrap();
188        assert_eq!(decoder.threshold_val(100), 99);
189    }
190}