Skip to main content

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

1//! Bit tree encoders from the LZMA SDK (`RangeCoderBitTree`).
2
3use super::range_coder::Encoder;
4use super::range_coder_bit::BitEncoder;
5
6/// A tree of binary probability models, addressed by the bits of a symbol.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct BitTreeEncoder {
9    models: Vec<BitEncoder>,
10    num_bit_levels: u32,
11}
12
13impl BitTreeEncoder {
14    /// Creates a tree with `num_bit_levels` levels.
15    pub fn new(num_bit_levels: u32) -> Self {
16        Self {
17            models: vec![BitEncoder::new(); 1usize << num_bit_levels],
18            num_bit_levels,
19        }
20    }
21
22    /// The number of levels of the tree.
23    pub fn num_bit_levels(&self) -> u32 {
24        self.num_bit_levels
25    }
26
27    /// Resets every model of the tree.
28    pub fn init(&mut self) {
29        for model in self.models.iter_mut().skip(1) {
30            model.init();
31        }
32    }
33
34    /// Encodes a symbol, most significant bit first.
35    pub fn encode(&mut self, range_encoder: &mut Encoder, symbol: u32) {
36        let mut m = 1u32;
37        let mut bit_index = self.num_bit_levels as i32;
38        while bit_index > 0 {
39            bit_index -= 1;
40            let bit = (symbol >> bit_index) & 1;
41            self.models[m as usize].encode(range_encoder, bit);
42            m = (m << 1) | bit;
43        }
44    }
45
46    /// Encodes a symbol, least significant bit first.
47    pub fn reverse_encode(&mut self, range_encoder: &mut Encoder, symbol: u32) {
48        let mut m = 1u32;
49        let mut symbol = symbol;
50        for _ in 0..self.num_bit_levels {
51            let bit = symbol & 1;
52            self.models[m as usize].encode(range_encoder, bit);
53            m = (m << 1) | bit;
54            symbol >>= 1;
55        }
56    }
57
58    /// The price of encoding a symbol, most significant bit first.
59    pub fn get_price(&self, symbol: u32) -> u32 {
60        let mut price = 0u32;
61        let mut m = 1u32;
62        let mut bit_index = self.num_bit_levels as i32;
63        while bit_index > 0 {
64            bit_index -= 1;
65            let bit = (symbol >> bit_index) & 1;
66            price += self.models[m as usize].get_price(bit);
67            m = (m << 1) + bit;
68        }
69        price
70    }
71
72    /// The price of encoding a symbol, least significant bit first.
73    pub fn reverse_get_price(&self, symbol: u32) -> u32 {
74        let mut price = 0u32;
75        let mut m = 1u32;
76        let mut symbol = symbol;
77        let mut i = self.num_bit_levels as i32;
78        while i > 0 {
79            let bit = symbol & 1;
80            symbol >>= 1;
81            price += self.models[m as usize].get_price(bit);
82            m = (m << 1) | bit;
83            i -= 1;
84        }
85        price
86    }
87
88    /// The price of encoding a symbol with an external tree, least significant
89    /// bit first.
90    pub fn reverse_get_price_at(
91        models: &[BitEncoder],
92        start_index: u32,
93        num_bit_levels: i32,
94        symbol: u32,
95    ) -> u32 {
96        let mut price = 0u32;
97        let mut m = 1u32;
98        let mut symbol = symbol;
99        let mut i = num_bit_levels;
100        while i > 0 {
101            let bit = symbol & 1;
102            symbol >>= 1;
103            price += models[start_index.wrapping_add(m) as usize].get_price(bit);
104            m = (m << 1) | bit;
105            i -= 1;
106        }
107        price
108    }
109
110    /// Encodes a symbol with an external tree, least significant bit first.
111    pub fn reverse_encode_at(
112        models: &mut [BitEncoder],
113        start_index: u32,
114        range_encoder: &mut Encoder,
115        num_bit_levels: i32,
116        symbol: u32,
117    ) {
118        let mut m = 1u32;
119        let mut symbol = symbol;
120        for _ in 0..num_bit_levels {
121            let bit = symbol & 1;
122            models[start_index.wrapping_add(m) as usize].encode(range_encoder, bit);
123            m = (m << 1) | bit;
124            symbol >>= 1;
125        }
126    }
127}