Skip to main content

libribzip2/block/symbol_statistics/
mod.rs

1use self::kmeansclustering::KMeansProblem;
2
3use super::zle::ZleSymbol;
4mod kmeansclustering;
5
6#[derive(Clone, Copy)]
7#[non_exhaustive]
8/// Strategy for using Huffman Tables
9///
10/// * [EncodingStrategy::Single] - every block of 900k uses a single huffman code table
11/// * [EncodingStrategy::BlockWise] - usage of code tables in 50 byte chunks is optimized using Lloyd's algorithm with given parameters
12pub enum EncodingStrategy {
13    BlockWise {
14        num_clusters: usize,
15        num_iterations: usize,
16    },
17    Single,
18}
19
20pub(crate) struct SinglePropabilityMap {
21    frequencies: Vec<usize>,
22    symbol_count: usize,
23}
24
25impl SinglePropabilityMap {
26    pub(crate) fn create(size: usize) -> Self {
27        SinglePropabilityMap {
28            frequencies: vec![0; size + 1],
29            symbol_count: 0,
30        }
31    }
32}
33
34impl SymbolReporter for SinglePropabilityMap {
35    fn report_symbol(&mut self, symbol: &ZleSymbol) {
36        self.symbol_count += 1;
37        match symbol {
38            ZleSymbol::RunA => {
39                self.frequencies[0] += 1;
40            }
41            ZleSymbol::RunB => {
42                self.frequencies[1] += 1;
43            }
44            ZleSymbol::Number(i) => {
45                self.frequencies[*i as usize + 1] += 1;
46            }
47        }
48    }
49
50    fn finalize(&mut self) -> ReportedSymbols {
51        let table = IntoFrequencyTable {
52            frequencies: self.frequencies.clone(),
53        };
54        ReportedSymbols {
55            reported_frequencies: vec![table.clone(), table],
56            selectors: vec![0; (self.symbol_count as f32 / 50.0).ceil() as usize],
57        }
58    }
59}
60
61pub(crate) struct BlockWisePropabilityMap {
62    current_frequencies: Vec<u8>,
63    pub(crate) maps: Vec<Vec<u8>>,
64    pub(crate) size: usize,
65    counter: usize,
66    num_iterations: usize,
67    num_clusters: usize,
68}
69
70impl BlockWisePropabilityMap {
71    pub(crate) fn create(size: usize, num_clusters: usize, num_iterations: usize) -> Self {
72        Self {
73            current_frequencies: vec![0; size + 1],
74            maps: vec![],
75            counter: 0,
76            size,
77            num_clusters,
78            num_iterations,
79        }
80    }
81}
82
83impl SymbolReporter for BlockWisePropabilityMap {
84    fn report_symbol(&mut self, symbol: &ZleSymbol) {
85        match symbol {
86            ZleSymbol::RunA => {
87                self.current_frequencies[0] += 1;
88            }
89            ZleSymbol::RunB => {
90                self.current_frequencies[1] += 1;
91            }
92            ZleSymbol::Number(i) => {
93                self.current_frequencies[*i as usize + 1] += 1;
94            }
95        }
96        self.counter += 1;
97        if self.counter >= 50 {
98            self.maps.push(std::mem::replace(
99                &mut self.current_frequencies,
100                vec![0; self.size + 1],
101            ));
102            self.counter = 0;
103        }
104    }
105
106    fn finalize(&mut self) -> ReportedSymbols {
107        self.maps.push(std::mem::replace(
108            &mut self.current_frequencies,
109            vec![0; self.size + 1],
110        ));
111
112        let p = KMeansProblem {
113            dimension: self.size + 1,
114            data: &self.maps,
115            num_iterations: self.num_iterations,
116            num_clusters: self.num_clusters,
117        };
118        let tables = p.solve();
119
120        ReportedSymbols {
121            reported_frequencies: tables
122                .means
123                .iter()
124                .map(|x| IntoFrequencyTable {
125                    frequencies: x.iter().map(|x| *x as usize).collect::<Vec<_>>(),
126                })
127                .collect::<Vec<_>>(),
128            selectors: tables.assignments,
129        }
130    }
131}
132
133pub(crate) trait SymbolReporter {
134    fn report_symbol(&mut self, symbol: &ZleSymbol);
135    fn finalize(&mut self) -> ReportedSymbols;
136}
137
138pub(crate) struct ReportedSymbols {
139    pub(crate) reported_frequencies: Vec<IntoFrequencyTable>,
140    pub(crate) selectors: Vec<u8>,
141}
142
143#[derive(Clone)]
144pub(crate) struct IntoFrequencyTable {
145    pub(crate) frequencies: Vec<usize>,
146}
147
148impl IntoFrequencyTable {
149    pub(crate) fn iterate(self) -> impl Iterator<Item = (ZleSymbol, usize)> {
150        self.frequencies
151            .into_iter()
152            .enumerate()
153            .map(|(symbol, frequency)| match symbol {
154                0 => (ZleSymbol::RunA, frequency),
155                1 => (ZleSymbol::RunB, frequency),
156                x => (ZleSymbol::Number((x - 1) as u8), frequency),
157            })
158    }
159}