Skip to main content

steamkit/util/lzma/compress/range_coder/
range_coder.rs

1//! The range encoder from the LZMA SDK (`RangeCoder`).
2//!
3//! Output goes into an owned `Vec<u8>`. `low` is a 64 bit accumulator whose
4//! upper half holds the carry; `range` is a u32.
5
6/// The binary range encoder.
7#[derive(Debug, Clone, Default)]
8pub struct Encoder {
9    /// The output buffer.
10    stream: Vec<u8>,
11
12    /// The low bound of the current interval.
13    pub low: u64,
14
15    /// The size of the current interval.
16    pub range: u32,
17
18    cache_size: u32,
19    cache: u8,
20
21    start_position: i64,
22}
23
24impl Encoder {
25    /// The value above which the range is renormalised.
26    pub const TOP_VALUE: u32 = 1 << 24;
27
28    /// Creates an uninitialised encoder.
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    /// Replaces the output buffer.
34    pub fn set_stream(&mut self, stream: Vec<u8>) {
35        self.stream = stream;
36    }
37
38    /// Takes the output buffer out of the encoder.
39    pub fn release_stream(&mut self) -> Vec<u8> {
40        std::mem::take(&mut self.stream)
41    }
42
43    /// The bytes written so far.
44    pub fn stream(&self) -> &[u8] {
45        &self.stream
46    }
47
48    /// The bytes written so far, mutably.
49    pub fn stream_mut(&mut self) -> &mut Vec<u8> {
50        &mut self.stream
51    }
52
53    /// Initialises the encoder for a new stream.
54    pub fn init(&mut self) {
55        self.start_position = self.stream.len() as i64;
56
57        self.low = 0;
58        self.range = 0xFFFF_FFFF;
59        self.cache_size = 1;
60        self.cache = 0;
61    }
62
63    /// Flushes the pending low bytes.
64    pub fn flush_data(&mut self) {
65        for _ in 0..5 {
66            self.shift_low();
67        }
68    }
69
70    /// Flushes the output stream.
71    ///
72    /// Does nothing, since a `Vec<u8>` has nothing to flush.
73    pub fn flush_stream(&mut self) {}
74
75    /// Closes the output stream.
76    ///
77    /// Drops the output buffer.
78    pub fn close_stream(&mut self) {
79        self.stream = Vec::new();
80    }
81
82    /// Encodes a symbol that occupies `size` of `total` probability units.
83    pub fn encode(&mut self, start: u32, size: u32, total: u32) {
84        let range_div = self.range / total;
85        self.range = range_div;
86        // u32 * u32 wraps before being widened to u64, same as the SDK.
87        self.low = self
88            .low
89            .wrapping_add(start.wrapping_mul(range_div) as u64);
90        self.range = self.range.wrapping_mul(size);
91        while self.range < Self::TOP_VALUE {
92            self.range <<= 8;
93            self.shift_low();
94        }
95    }
96
97    /// Shifts the low accumulator by one byte, emitting the pending bytes.
98    pub fn shift_low(&mut self) {
99        let low32 = self.low as u32;
100        if low32 < 0xFF00_0000 || (self.low >> 32) == 1 {
101            let mut temp = self.cache;
102            loop {
103                self.stream
104                    .push((temp as u32 + (self.low >> 32) as u32) as u8);
105                temp = 0xFF;
106                self.cache_size -= 1;
107                if self.cache_size == 0 {
108                    break;
109                }
110            }
111            self.cache = (low32 >> 24) as u8;
112        }
113        self.cache_size += 1;
114        self.low = (low32 << 8) as u64;
115    }
116
117    /// Encodes `num_total_bits` bits of `v` without a probability model.
118    pub fn encode_direct_bits(&mut self, v: u32, num_total_bits: i32) {
119        let mut i = num_total_bits - 1;
120        while i >= 0 {
121            self.range >>= 1;
122            if ((v >> i) & 1) == 1 {
123                self.low = self.low.wrapping_add(self.range as u64);
124            }
125            if self.range < Self::TOP_VALUE {
126                self.range <<= 8;
127                self.shift_low();
128            }
129            i -= 1;
130        }
131    }
132
133    /// Encodes a single bit using a probability model of `num_total_bits` bits.
134    pub fn encode_bit(&mut self, size0: u32, num_total_bits: i32, symbol: u32) {
135        let new_bound = (self.range >> num_total_bits) * size0;
136        if symbol == 0 {
137            self.range = new_bound;
138        } else {
139            self.low = self.low.wrapping_add(new_bound as u64);
140            self.range -= new_bound;
141        }
142        while self.range < Self::TOP_VALUE {
143            self.range <<= 8;
144            self.shift_low();
145        }
146    }
147
148    /// The number of bytes produced so far, plus the four bytes that the
149    /// decoder needs to finish the stream.
150    pub fn get_processed_size_add(&self) -> i64 {
151        self.cache_size as i64 + (self.stream.len() as i64 - self.start_position) + 4
152    }
153}