preflate_rs/
statistical_codec.rs

1/*---------------------------------------------------------------------------------------------
2 *  Copyright (c) Microsoft Corporation. All rights reserved.
3 *  Licensed under the Apache License, Version 2.0. See LICENSE.txt in the project root for license information.
4 *  This software incorporates material from third parties. See NOTICE.txt for details.
5 *--------------------------------------------------------------------------------------------*/
6
7use crate::cabac_codec::{decode_difference, encode_difference};
8
9/// correction indictions, which are followed by a 16 bit value
10#[derive(Copy, Clone, Debug, Eq, PartialEq)]
11pub enum CodecCorrection {
12    TokenCount,
13    BlockTypeCorrection,
14    LenCorrection,
15    DistOnlyCorrection,
16    DistAfterLenCorrection,
17    TreeCodeBitLengthCorrection,
18    LDTypeCorrection,
19    RepeatCountCorrection,
20    LDBitLengthCorrection,
21
22    TreeCodeCountCorrection,
23    LiteralCountCorrection,
24    DistanceCountCorrection,
25    UncompressBlockLenCorrection,
26
27    Last,
28    EndOfChunk,
29    LiteralPredictionWrong,
30    ReferencePredictionWrong,
31    MAX,
32}
33
34pub trait PredictionEncoder {
35    fn encode_correction(&mut self, action: CodecCorrection, value: u32);
36
37    fn encode_verify_state(&mut self, message: &'static str, checksum: u64);
38
39    fn finish(&mut self);
40
41    fn encode_correction_diff(
42        &mut self,
43        action: CodecCorrection,
44        actual_value: u32,
45        predicted_value: u32,
46    ) {
47        self.encode_correction(action, encode_difference(predicted_value, actual_value));
48    }
49
50    fn encode_misprediction(&mut self, action: CodecCorrection, actual_value: bool) {
51        self.encode_correction(action, actual_value as u32);
52    }
53
54    /// if a bool is not equal, then we encode a 1, otherwise a 0
55    fn encode_correction_bool(
56        &mut self,
57        action: CodecCorrection,
58        actual_value: bool,
59        predicted_value: bool,
60    ) {
61        self.encode_correction(action, (actual_value != predicted_value) as u32)
62    }
63}
64
65pub trait PredictionDecoder {
66    fn decode_correction(&mut self, correction: CodecCorrection) -> u32;
67    fn decode_verify_state(&mut self, message: &'static str, checksum: u64);
68
69    #[inline]
70    fn decode_correction_diff(&mut self, correction: CodecCorrection, predicted_value: u32) -> u32 {
71        let actual_value = self.decode_correction(correction);
72        decode_difference(predicted_value, actual_value)
73    }
74
75    #[inline]
76    fn decode_misprediction(&mut self, correction: CodecCorrection) -> bool {
77        self.decode_correction(correction) != 0
78    }
79
80    /// if we encoded a 1 then swap the predicted value
81    fn decode_correction_bool(
82        &mut self,
83        correction: CodecCorrection,
84        predicted_value: bool,
85    ) -> bool {
86        predicted_value ^ (self.decode_correction(correction) != 0)
87    }
88}
89
90#[derive(Copy, Clone, Debug, Eq, PartialEq)]
91pub enum CodecAction {
92    Correction(CodecCorrection, u32),
93    VerifyState(&'static str, u64),
94}
95
96#[derive(Default)]
97pub struct CountNonDefaultActions {
98    pub total_non_default: u32,
99
100    pub corrections_count: [u32; CodecCorrection::MAX as usize],
101}
102
103impl CountNonDefaultActions {
104    pub fn record_correction(&mut self, correction: CodecCorrection, value: u32) {
105        if value != 0 {
106            self.corrections_count[correction as usize] += 1;
107            self.total_non_default += 1;
108        }
109    }
110
111    pub fn print(&self) {
112        use CodecCorrection::*;
113
114        let corr = [
115            TokenCount,
116            BlockTypeCorrection,
117            LenCorrection,
118            DistOnlyCorrection,
119            DistAfterLenCorrection,
120            TreeCodeBitLengthCorrection,
121            LDTypeCorrection,
122            RepeatCountCorrection,
123            LDBitLengthCorrection,
124            TreeCodeCountCorrection,
125            LiteralCountCorrection,
126            DistanceCountCorrection,
127            UncompressBlockLenCorrection,
128            Last,
129            EndOfChunk,
130            LiteralPredictionWrong,
131            ReferencePredictionWrong,
132        ];
133
134        assert_eq!(
135            corr.len(),
136            CodecCorrection::MAX as usize,
137            "need to update array if you add an enum"
138        );
139
140        for i in corr {
141            if self.corrections_count[i as usize] != 0 {
142                println!("{:?}: {}", i, self.corrections_count[i as usize]);
143            }
144        }
145    }
146}
147
148pub struct VerifyPredictionDecoder {
149    actions: Vec<CodecAction>,
150    index: usize,
151}
152
153#[derive(Default)]
154pub struct VerifyPredictionEncoder {
155    actions: Vec<CodecAction>,
156    count: CountNonDefaultActions,
157}
158
159// used for testing mostly
160#[allow(dead_code)]
161impl VerifyPredictionEncoder {
162    pub fn new() -> Self {
163        Self {
164            actions: Vec::new(),
165            count: CountNonDefaultActions::default(),
166        }
167    }
168
169    pub fn actions(&self) -> Vec<CodecAction> {
170        self.actions.clone()
171    }
172
173    pub fn print(&self) {
174        self.count.print();
175    }
176
177    pub fn count_nondefault_actions(&self) -> usize {
178        self.count.total_non_default as usize
179    }
180}
181
182impl PredictionEncoder for VerifyPredictionEncoder {
183    fn encode_verify_state(&mut self, message: &'static str, checksum: u64) {
184        self.actions
185            .push(CodecAction::VerifyState(message, checksum));
186    }
187
188    fn encode_correction(&mut self, action: CodecCorrection, value: u32) {
189        self.actions.push(CodecAction::Correction(action, value));
190        self.count.record_correction(action, value);
191    }
192    fn finish(&mut self) {}
193}
194
195// used for testing mostly
196#[allow(dead_code)]
197impl VerifyPredictionDecoder {
198    pub fn new(actions: Vec<CodecAction>) -> Self {
199        Self { actions, index: 0 }
200    }
201
202    fn pop(&mut self) -> Option<CodecAction> {
203        if self.index >= self.actions.len() {
204            None
205        } else {
206            self.index += 1;
207            Some(self.actions[self.index - 1])
208        }
209    }
210}
211
212impl PredictionDecoder for VerifyPredictionDecoder {
213    fn decode_verify_state(&mut self, message: &'static str, checksum: u64) {
214        let x = self.pop().unwrap();
215        assert_eq!(
216            x,
217            CodecAction::VerifyState(message, checksum),
218            "mismatch {} (left encode, right decode)",
219            self.index
220        );
221    }
222
223    fn decode_correction(&mut self, correction: CodecCorrection) -> u32 {
224        let x = self.pop().unwrap();
225        match x {
226            CodecAction::Correction(c, value) => {
227                assert_eq!(correction, c);
228                return value;
229            }
230            CodecAction::VerifyState(s, _h) => {
231                panic!("found VerifyState {}, expected {:?}", s, correction);
232            }
233        }
234    }
235}
236
237#[cfg(test)]
238pub fn drive_encoder<T: PredictionEncoder>(encoder: &mut T, actions: &[CodecAction]) {
239    for action in actions {
240        match action {
241            &CodecAction::Correction(correction, value) => {
242                encoder.encode_correction(correction, value);
243            }
244            &CodecAction::VerifyState(message, checksum) => {
245                encoder.encode_verify_state(message, checksum);
246            }
247        }
248    }
249}
250
251#[cfg(test)]
252pub fn verify_decoder<T: PredictionDecoder>(decoder: &mut T, actions: &[CodecAction]) {
253    for action in actions {
254        match action {
255            &CodecAction::Correction(correction, value) => {
256                let x = decoder.decode_correction(correction);
257                assert_eq!(x, value);
258            }
259            &CodecAction::VerifyState(message, checksum) => {
260                decoder.decode_verify_state(message, checksum);
261            }
262        }
263    }
264}
265
266// used by tests to ensure that perfect prediction is performed for data
267// that we know should be encoded without any mispredictions or corrections
268#[cfg(test)]
269#[derive(Default)]
270pub struct AssertDefaultOnlyEncoder {}
271
272#[cfg(test)]
273impl PredictionEncoder for AssertDefaultOnlyEncoder {
274    fn encode_correction(&mut self, action: CodecCorrection, value: u32) {
275        assert_eq!(0, value, "unexpected correction {:?}", action);
276    }
277
278    fn encode_verify_state(&mut self, _message: &'static str, _checksum: u64) {}
279
280    fn finish(&mut self) {}
281}
282
283#[cfg(test)]
284#[derive(Default)]
285pub struct AssertDefaultOnlyDecoder {}
286
287#[cfg(test)]
288impl PredictionDecoder for AssertDefaultOnlyDecoder {
289    fn decode_correction(&mut self, _correction: CodecCorrection) -> u32 {
290        0
291    }
292
293    fn decode_verify_state(&mut self, _message: &'static str, _checksum: u64) {}
294}
295
296/// This implements a prediction encoder that tees the input to two different
297/// encoders. This allows us to verify that the behavior of two encoders is the same
298impl<A, B> PredictionEncoder for (A, B)
299where
300    A: PredictionEncoder,
301    B: PredictionEncoder,
302{
303    fn encode_verify_state(&mut self, message: &'static str, checksum: u64) {
304        self.0.encode_verify_state(message, checksum);
305        self.1.encode_verify_state(message, checksum);
306    }
307
308    fn encode_correction(&mut self, action: CodecCorrection, value: u32) {
309        self.0.encode_correction(action, value);
310        self.1.encode_correction(action, value);
311    }
312
313    fn finish(&mut self) {
314        self.0.finish();
315        self.1.finish();
316    }
317}
318
319/// Implement the same for decoders, where we verify that the output
320/// is identical for both decoders
321impl<A, B> PredictionDecoder for (A, B)
322where
323    A: PredictionDecoder,
324    B: PredictionDecoder,
325{
326    fn decode_correction(&mut self, correction: CodecCorrection) -> u32 {
327        let a = self.0.decode_correction(correction);
328        let b = self.1.decode_correction(correction);
329        assert_eq!(a, b);
330        a
331    }
332
333    fn decode_verify_state(&mut self, message: &'static str, checksum: u64) {
334        self.0.decode_verify_state(message, checksum);
335        self.1.decode_verify_state(message, checksum);
336    }
337}
preflate_rs/statistical_codec.rs

preflate_rs/
statistical_codec.rs
