arithmetic_coding_adder_dep/
decoder.rs

1//! The [`Decoder`] half of the arithmetic coding library.
2
3use std::marker::PhantomData;
4use std::{io, ops::Range};
5
6use bitstream_io::BitRead;
7
8use crate::{BitStore, Model};
9
10// this algorithm is derived from this article - https://marknelson.us/posts/2014/10/19/data-compression-with-arithmetic-coding.html
11
12/// An arithmetic decoder
13///
14/// An arithmetic decoder converts a stream of bytes into a stream of some
15/// output symbol, using a predictive [`Model`].
16#[derive(Debug)]
17pub struct Decoder<M, R>
18where
19    M: Model,
20    R: BitRead,
21{
22    /// The model used to predict the next symbol
23    pub model: M,
24    state: State<M::B, R>,
25}
26
27trait BitReadExt {
28    fn next_bit(&mut self) -> io::Result<Option<bool>>;
29}
30
31impl<R: BitRead> BitReadExt for R {
32    fn next_bit(&mut self) -> io::Result<Option<bool>> {
33        match self.read_bit() {
34            Ok(bit) => Ok(Some(bit)),
35            Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => Ok(None),
36            Err(e) => Err(e),
37        }
38    }
39}
40
41impl<M, R> Decoder<M, R>
42where
43    M: Model,
44    R: BitRead,
45{
46    /// Construct a new [`Decoder`]
47    ///
48    /// The 'precision' of the encoder is maximised, based on the number of bits
49    /// needed to represent the [`Model::denominator`]. 'precision' bits is
50    /// equal to [`u32::BITS`] - [`Model::denominator`] bits.
51    ///
52    /// # Panics
53    ///
54    /// The calculation of the number of bits used for 'precision' is subject to
55    /// the following constraints:
56    ///
57    /// - The total available bits is [`u32::BITS`]
58    /// - The precision must use at least 2 more bits than that needed to
59    ///   represent [`Model::denominator`]
60    ///
61    /// If these constraints cannot be satisfied this method will panic in debug
62    /// builds
63    pub fn new(model: M) -> Self {
64        let frequency_bits = model.max_denominator().log2() + 1;
65        let precision = M::B::BITS - frequency_bits;
66
67        Self::with_precision(model, precision)
68    }
69
70    /// Construct a new [`Decoder`] with a custom precision
71    ///
72    /// # Panics
73    ///
74    /// The calculation of the number of bits used for 'precision' is subject to
75    /// the following constraints:
76    ///
77    /// - The total available bits is [`BitStore::BITS`]
78    /// - The precision must use at least 2 more bits than that needed to
79    ///   represent [`Model::denominator`]
80    ///
81    /// If these constraints cannot be satisfied this method will panic in debug
82    /// builds
83    pub fn with_precision(model: M, precision: u32) -> Self {
84        let frequency_bits = model.max_denominator().log2() + 1;
85        debug_assert!(
86            (precision >= (frequency_bits + 2)),
87            "not enough bits of precision to prevent overflow/underflow",
88        );
89        debug_assert!(
90            (frequency_bits + precision) <= M::B::BITS,
91            "not enough bits in BitStore to support the required precision",
92        );
93
94        let state = State::new(precision);
95
96        Self { model, state }
97    }
98
99    /// todo
100    pub const fn with_state(state: State<M::B, R>, model: M) -> Self {
101        Self { model, state }
102    }
103
104    /// Return an iterator over the decoded symbols.
105    ///
106    /// The iterator will continue returning symbols until EOF is reached
107    pub fn decode_all<'a>(&'a mut self, input: &'a mut R) -> DecodeIter<'a, M, R> {
108        DecodeIter {
109            decoder: self,
110            input,
111        }
112    }
113
114    /// Read the next symbol from the stream of bits
115    ///
116    /// This method will return `Ok(None)` when EOF is reached.
117    ///
118    /// # Errors
119    ///
120    /// This method can fail if the underlying [`BitRead`] cannot be read from.
121    pub fn decode(&mut self, input: &mut R) -> io::Result<Option<M::Symbol>> {
122        self.state.initialise(input)?;
123
124        let denominator = self.model.denominator();
125        debug_assert!(
126            denominator <= self.model.max_denominator(),
127            "denominator is greater than maximum!"
128        );
129        let value = self.state.value(denominator);
130        let symbol = self.model.symbol(value);
131
132        let p = self
133            .model
134            .probability(symbol.as_ref())
135            .expect("this should not be able to fail. Check the implementation of the model.");
136
137        self.state.scale(p, denominator, input)?;
138        self.model.update(symbol.as_ref());
139
140        Ok(symbol)
141    }
142
143    /// Reuse the internal state of the Decoder with a new model.
144    ///
145    /// Allows for chaining multiple sequences of symbols from a single stream
146    /// of bits
147    pub fn chain<X>(self, model: X) -> Decoder<X, R>
148    where
149        X: Model<B = M::B>,
150    {
151        Decoder {
152            model,
153            state: self.state,
154        }
155    }
156
157    /// todo
158    pub fn into_inner(self) -> (M, State<M::B, R>) {
159        (self.model, self.state)
160    }
161}
162
163/// The iterator returned by the [`Model::decode_all`] method
164#[allow(missing_debug_implementations)]
165pub struct DecodeIter<'a, M, R>
166where
167    M: Model,
168    R: BitRead,
169{
170    decoder: &'a mut Decoder<M, R>,
171    input: &'a mut R,
172}
173
174impl<M, R> Iterator for DecodeIter<'_, M, R>
175where
176    M: Model,
177    R: BitRead,
178{
179    type Item = io::Result<M::Symbol>;
180
181    fn next(&mut self) -> Option<Self::Item> {
182        self.decoder.decode(self.input).transpose()
183    }
184}
185
186/// A convenience struct which stores the internal state of an [`Decoder`].
187#[derive(Debug)]
188pub struct State<B, R>
189where
190    B: BitStore,
191    R: BitRead,
192{
193    precision: u32,
194    low: B,
195    high: B,
196    _marker: PhantomData<R>,
197    x: B,
198    uninitialised: bool,
199}
200
201impl<B, R> State<B, R>
202where
203    B: BitStore,
204    R: BitRead,
205{
206    /// todo
207    #[must_use]
208    pub fn new(precision: u32) -> Self {
209        let low = B::ZERO;
210        let high = B::ONE << precision;
211        let x = B::ZERO;
212
213        Self {
214            precision,
215            low,
216            high,
217            _marker: PhantomData,
218            x,
219            uninitialised: true,
220        }
221    }
222
223    fn half(&self) -> B {
224        B::ONE << (self.precision - 1)
225    }
226
227    fn quarter(&self) -> B {
228        B::ONE << (self.precision - 2)
229    }
230
231    fn three_quarter(&self) -> B {
232        self.half() + self.quarter()
233    }
234
235    fn normalise(&mut self, input: &mut R) -> io::Result<()> {
236        while self.high < self.half() || self.low >= self.half() {
237            if self.high < self.half() {
238                self.high <<= 1;
239                self.low <<= 1;
240                self.x <<= 1;
241            } else {
242                // self.low >= self.half()
243                self.low = (self.low - self.half()) << 1;
244                self.high = (self.high - self.half()) << 1;
245                self.x = (self.x - self.half()) << 1;
246            }
247
248            if let Some(true) = input.next_bit()? {
249                self.x += B::ONE;
250            }
251        }
252
253        while self.low >= self.quarter() && self.high < (self.three_quarter()) {
254            self.low = (self.low - self.quarter()) << 1;
255            self.high = (self.high - self.quarter()) << 1;
256            self.x = (self.x - self.quarter()) << 1;
257
258            if input.next_bit()? == Some(true) {
259                self.x += B::ONE;
260            }
261        }
262
263        Ok(())
264    }
265
266    fn scale(&mut self, p: Range<B>, denominator: B, input: &mut R) -> io::Result<()> {
267        let range = self.high - self.low + B::ONE;
268
269        self.high = self.low + (range * p.end) / denominator - B::ONE;
270        self.low += (range * p.start) / denominator;
271
272        self.normalise(input)
273    }
274
275    fn value(&self, denominator: B) -> B {
276        let range = self.high - self.low + B::ONE;
277        ((self.x - self.low + B::ONE) * denominator - B::ONE) / range
278    }
279
280    fn fill(&mut self, input: &mut R) -> io::Result<()> {
281        for _ in 0..self.precision {
282            self.x <<= 1;
283            if input.next_bit()? == Some(true) {
284                self.x += B::ONE;
285            }
286        }
287        Ok(())
288    }
289
290    fn initialise(&mut self, input: &mut R) -> io::Result<()> {
291        if self.uninitialised {
292            self.fill(input)?;
293            self.uninitialised = false;
294        }
295        Ok(())
296    }
297}
arithmetic_coding_adder_dep/decoder.rs

arithmetic_coding_adder_dep/
decoder.rs
