Skip to main content

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

1//! Binary probability models from the LZMA SDK (`RangeCoderBit`).
2//!
3//! Probabilities are updated with the usual 7-Zip formula
4//! (`prob += (2048 - prob) >> 5` or `prob -= prob >> 5`). The price table
5//! matches the SDK exactly, since the encoder's choices depend on it.
6
7use super::range_coder::Encoder;
8
9/// Static table of bit prices.
10const PROB_PRICES: [u32; 512] = build_prob_prices();
11
12const NUM_BITS: i32 = 9;
13
14const fn build_prob_prices() -> [u32; 512] {
15    let mut table = [0u32; 512];
16    let mut i = NUM_BITS - 1;
17    while i >= 0 {
18        let start = 1u32 << (NUM_BITS - i - 1);
19        let end = 1u32 << (NUM_BITS - i);
20        let mut j = start;
21        while j < end {
22            table[j as usize] = ((i as u32) << 6) + (((end - j) << 6) >> (NUM_BITS - i - 1));
23            j += 1;
24        }
25        i -= 1;
26    }
27    table
28}
29
30/// A single binary probability model.
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
32pub struct BitEncoder {
33    prob: u32,
34}
35
36impl BitEncoder {
37    /// The number of bits used to represent a probability.
38    pub const NUM_BIT_MODEL_TOTAL_BITS: i32 = 11;
39
40    /// The total number of probability units.
41    pub const BIT_MODEL_TOTAL: u32 = 1 << Self::NUM_BIT_MODEL_TOTAL_BITS;
42
43    /// The number of bits the probability is shifted by on every update.
44    pub const NUM_MOVE_BITS: i32 = 5;
45
46    /// The number of bits the probability is reduced by when pricing.
47    pub const NUM_MOVE_REDUCING_BITS: i32 = 2;
48
49    /// The shift applied to the price of a bit.
50    pub const NUM_BIT_PRICE_SHIFT_BITS: i32 = 6;
51
52    /// Creates an uninitialised model (probability zero).
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    /// Resets the probability to its initial value.
58    pub fn init(&mut self) {
59        self.prob = Self::BIT_MODEL_TOTAL >> 1;
60    }
61
62    /// The current probability.
63    pub fn prob(&self) -> u32 {
64        self.prob
65    }
66
67    /// Updates the probability for the given symbol.
68    pub fn update_model(&mut self, symbol: u32) {
69        if symbol == 0 {
70            self.prob += (Self::BIT_MODEL_TOTAL - self.prob) >> Self::NUM_MOVE_BITS;
71        } else {
72            self.prob -= self.prob >> Self::NUM_MOVE_BITS;
73        }
74    }
75
76    /// Encodes one bit with this model and updates it.
77    pub fn encode(&mut self, encoder: &mut Encoder, symbol: u32) {
78        let new_bound = (encoder.range >> Self::NUM_BIT_MODEL_TOTAL_BITS) * self.prob;
79        if symbol == 0 {
80            encoder.range = new_bound;
81            self.prob += (Self::BIT_MODEL_TOTAL - self.prob) >> Self::NUM_MOVE_BITS;
82        } else {
83            encoder.low = encoder.low.wrapping_add(new_bound as u64);
84            encoder.range -= new_bound;
85            self.prob -= self.prob >> Self::NUM_MOVE_BITS;
86        }
87        if encoder.range < Encoder::TOP_VALUE {
88            encoder.range <<= 8;
89            encoder.shift_low();
90        }
91    }
92
93    /// The price of encoding `symbol`.
94    pub fn get_price(&self, symbol: u32) -> u32 {
95        let index = (self.prob.wrapping_sub(symbol) ^ ((symbol as i32).wrapping_neg() as u32))
96            & (Self::BIT_MODEL_TOTAL - 1);
97        PROB_PRICES[(index >> Self::NUM_MOVE_REDUCING_BITS) as usize]
98    }
99
100    /// The price of encoding a zero bit.
101    pub fn get_price0(&self) -> u32 {
102        PROB_PRICES[(self.prob >> Self::NUM_MOVE_REDUCING_BITS) as usize]
103    }
104
105    /// The price of encoding a one bit.
106    pub fn get_price1(&self) -> u32 {
107        PROB_PRICES[((Self::BIT_MODEL_TOTAL - self.prob) >> Self::NUM_MOVE_REDUCING_BITS) as usize]
108    }
109}