scuffle_expgolomb/
lib.rs

1//! A set of helper functions to encode and decode exponential-golomb values.
2//!
3//! This crate extends upon the [`BitReader`] and [`BitWriter`] from the
4//! [`scuffle-bytes-util`](scuffle_bytes_util) crate to provide functionality
5//! for reading and writing Exp-Golomb encoded numbers.
6//!
7//! ```rust
8//! # fn test() -> std::io::Result<()> {
9//! use scuffle_expgolomb::{BitReaderExpGolombExt, BitWriterExpGolombExt};
10//! use scuffle_bytes_util::{BitReader, BitWriter};
11//!
12//! let mut bit_writer = BitWriter::default();
13//! bit_writer.write_exp_golomb(0)?;
14//! bit_writer.write_exp_golomb(1)?;
15//! bit_writer.write_exp_golomb(2)?;
16//!
17//! let data: Vec<u8> = bit_writer.finish()?;
18//!
19//! let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
20//!
21//! let result = bit_reader.read_exp_golomb()?;
22//! assert_eq!(result, 0);
23//!
24//! let result = bit_reader.read_exp_golomb()?;
25//! assert_eq!(result, 1);
26//!
27//! let result = bit_reader.read_exp_golomb()?;
28//! assert_eq!(result, 2);
29//! # Ok(())
30//! # }
31//! # test().expect("failed to run test");
32//! ```
33//!
34//! ## License
35//!
36//! This project is licensed under the [MIT](./LICENSE.MIT) or
37//! [Apache-2.0](./LICENSE.Apache-2.0) license. You can choose between one of
38//! them if you use this work.
39//!
40//! `SPDX-License-Identifier: MIT OR Apache-2.0`
41#![cfg_attr(all(coverage_nightly, test), feature(coverage_attribute))]
42
43use std::io;
44
45use scuffle_bytes_util::{BitReader, BitWriter};
46
47/// Extension trait for reading Exp-Golomb encoded numbers from a bit reader
48///
49/// See: <https://en.wikipedia.org/wiki/Exponential-Golomb_coding>
50///
51/// - [`BitReader`]
52pub trait BitReaderExpGolombExt {
53    /// Reads an Exp-Golomb encoded number
54    fn read_exp_golomb(&mut self) -> io::Result<u64>;
55
56    /// Reads a signed Exp-Golomb encoded number
57    fn read_signed_exp_golomb(&mut self) -> io::Result<i64> {
58        let exp_glob = self.read_exp_golomb()?;
59
60        if exp_glob % 2 == 0 {
61            Ok(-((exp_glob / 2) as i64))
62        } else {
63            Ok((exp_glob / 2) as i64 + 1)
64        }
65    }
66}
67
68impl<R: io::Read> BitReaderExpGolombExt for BitReader<R> {
69    fn read_exp_golomb(&mut self) -> io::Result<u64> {
70        let mut leading_zeros = 0;
71        while !self.read_bit()? {
72            leading_zeros += 1;
73        }
74
75        let mut result = 1;
76        for _ in 0..leading_zeros {
77            result <<= 1;
78            result |= self.read_bit()? as u64;
79        }
80
81        Ok(result - 1)
82    }
83}
84
85/// Extension trait for writing Exp-Golomb encoded numbers to a bit writer
86///
87/// See: <https://en.wikipedia.org/wiki/Exponential-Golomb_coding>
88///
89/// - [`BitWriter`]
90pub trait BitWriterExpGolombExt {
91    /// Writes an Exp-Golomb encoded number
92    fn write_exp_golomb(&mut self, input: u64) -> io::Result<()>;
93
94    /// Writes a signed Exp-Golomb encoded number
95    fn write_signed_exp_golomb(&mut self, number: i64) -> io::Result<()> {
96        let number = if number <= 0 {
97            -number as u64 * 2
98        } else {
99            number as u64 * 2 - 1
100        };
101
102        self.write_exp_golomb(number)
103    }
104}
105
106impl<W: io::Write> BitWriterExpGolombExt for BitWriter<W> {
107    fn write_exp_golomb(&mut self, input: u64) -> io::Result<()> {
108        let mut number = input + 1;
109        let mut leading_zeros = 0;
110        while number > 1 {
111            number >>= 1;
112            leading_zeros += 1;
113        }
114
115        for _ in 0..leading_zeros {
116            self.write_bit(false)?;
117        }
118
119        self.write_bits(input + 1, leading_zeros + 1)?;
120
121        Ok(())
122    }
123}
124
125#[cfg(test)]
126#[cfg_attr(all(test, coverage_nightly), coverage(off))]
127mod tests {
128    use bytes::Buf;
129    use scuffle_bytes_util::{BitReader, BitWriter};
130
131    use crate::{BitReaderExpGolombExt, BitWriterExpGolombExt};
132
133    pub fn get_remaining_bits(reader: &BitReader<std::io::Cursor<Vec<u8>>>) -> usize {
134        let remaining = reader.get_ref().remaining();
135
136        if reader.is_aligned() {
137            remaining * 8
138        } else {
139            remaining * 8 + (8 - reader.bit_pos() as usize)
140        }
141    }
142
143    #[test]
144    fn test_exp_glob_decode() {
145        let mut bit_writer = BitWriter::<Vec<u8>>::default();
146
147        bit_writer.write_bits(0b1, 1).unwrap(); // 0
148        bit_writer.write_bits(0b010, 3).unwrap(); // 1
149        bit_writer.write_bits(0b011, 3).unwrap(); // 2
150        bit_writer.write_bits(0b00100, 5).unwrap(); // 3
151        bit_writer.write_bits(0b00101, 5).unwrap(); // 4
152        bit_writer.write_bits(0b00110, 5).unwrap(); // 5
153        bit_writer.write_bits(0b00111, 5).unwrap(); // 6
154
155        let data = bit_writer.finish().unwrap();
156
157        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
158
159        let remaining_bits = get_remaining_bits(&bit_reader);
160
161        let result = bit_reader.read_exp_golomb().unwrap();
162        assert_eq!(result, 0);
163        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
164
165        let result = bit_reader.read_exp_golomb().unwrap();
166        assert_eq!(result, 1);
167        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
168
169        let result = bit_reader.read_exp_golomb().unwrap();
170        assert_eq!(result, 2);
171        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
172
173        let result = bit_reader.read_exp_golomb().unwrap();
174        assert_eq!(result, 3);
175        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
176
177        let result = bit_reader.read_exp_golomb().unwrap();
178        assert_eq!(result, 4);
179        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
180
181        let result = bit_reader.read_exp_golomb().unwrap();
182        assert_eq!(result, 5);
183        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
184
185        let result = bit_reader.read_exp_golomb().unwrap();
186        assert_eq!(result, 6);
187        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
188    }
189
190    #[test]
191    fn test_signed_exp_glob_decode() {
192        let mut bit_writer = BitWriter::<Vec<u8>>::default();
193
194        bit_writer.write_bits(0b1, 1).unwrap(); // 0
195        bit_writer.write_bits(0b010, 3).unwrap(); // 1
196        bit_writer.write_bits(0b011, 3).unwrap(); // -1
197        bit_writer.write_bits(0b00100, 5).unwrap(); // 2
198        bit_writer.write_bits(0b00101, 5).unwrap(); // -2
199        bit_writer.write_bits(0b00110, 5).unwrap(); // 3
200        bit_writer.write_bits(0b00111, 5).unwrap(); // -3
201
202        let data = bit_writer.finish().unwrap();
203
204        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
205
206        let remaining_bits = get_remaining_bits(&bit_reader);
207
208        let result = bit_reader.read_signed_exp_golomb().unwrap();
209        assert_eq!(result, 0);
210        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
211
212        let result = bit_reader.read_signed_exp_golomb().unwrap();
213        assert_eq!(result, 1);
214        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
215
216        let result = bit_reader.read_signed_exp_golomb().unwrap();
217        assert_eq!(result, -1);
218        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
219
220        let result = bit_reader.read_signed_exp_golomb().unwrap();
221        assert_eq!(result, 2);
222        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
223
224        let result = bit_reader.read_signed_exp_golomb().unwrap();
225        assert_eq!(result, -2);
226        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
227
228        let result = bit_reader.read_signed_exp_golomb().unwrap();
229        assert_eq!(result, 3);
230        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
231
232        let result = bit_reader.read_signed_exp_golomb().unwrap();
233        assert_eq!(result, -3);
234        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
235    }
236
237    #[test]
238    fn test_exp_glob_encode() {
239        let mut bit_writer = BitWriter::<Vec<u8>>::default();
240
241        bit_writer.write_exp_golomb(0).unwrap();
242        bit_writer.write_exp_golomb(1).unwrap();
243        bit_writer.write_exp_golomb(2).unwrap();
244        bit_writer.write_exp_golomb(3).unwrap();
245        bit_writer.write_exp_golomb(4).unwrap();
246        bit_writer.write_exp_golomb(5).unwrap();
247        bit_writer.write_exp_golomb(6).unwrap();
248        bit_writer.write_exp_golomb(u64::MAX - 1).unwrap();
249
250        let data = bit_writer.finish().unwrap();
251
252        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
253
254        let remaining_bits = get_remaining_bits(&bit_reader);
255
256        let result = bit_reader.read_exp_golomb().unwrap();
257        assert_eq!(result, 0);
258        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
259
260        let result = bit_reader.read_exp_golomb().unwrap();
261        assert_eq!(result, 1);
262        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
263
264        let result = bit_reader.read_exp_golomb().unwrap();
265        assert_eq!(result, 2);
266        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
267
268        let result = bit_reader.read_exp_golomb().unwrap();
269        assert_eq!(result, 3);
270        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
271
272        let result = bit_reader.read_exp_golomb().unwrap();
273        assert_eq!(result, 4);
274        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
275
276        let result = bit_reader.read_exp_golomb().unwrap();
277        assert_eq!(result, 5);
278        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
279
280        let result = bit_reader.read_exp_golomb().unwrap();
281        assert_eq!(result, 6);
282        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
283
284        let result = bit_reader.read_exp_golomb().unwrap();
285        assert_eq!(result, u64::MAX - 1);
286        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 154);
287    }
288
289    #[test]
290    fn test_signed_exp_glob_encode() {
291        let mut bit_writer = BitWriter::<Vec<u8>>::default();
292
293        bit_writer.write_signed_exp_golomb(0).unwrap();
294        bit_writer.write_signed_exp_golomb(1).unwrap();
295        bit_writer.write_signed_exp_golomb(-1).unwrap();
296        bit_writer.write_signed_exp_golomb(2).unwrap();
297        bit_writer.write_signed_exp_golomb(-2).unwrap();
298        bit_writer.write_signed_exp_golomb(3).unwrap();
299        bit_writer.write_signed_exp_golomb(-3).unwrap();
300        bit_writer.write_signed_exp_golomb(i64::MAX).unwrap();
301
302        let data = bit_writer.finish().unwrap();
303
304        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
305
306        let remaining_bits = get_remaining_bits(&bit_reader);
307
308        let result = bit_reader.read_signed_exp_golomb().unwrap();
309        assert_eq!(result, 0);
310        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
311
312        let result = bit_reader.read_signed_exp_golomb().unwrap();
313        assert_eq!(result, 1);
314        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
315
316        let result = bit_reader.read_signed_exp_golomb().unwrap();
317        assert_eq!(result, -1);
318        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
319
320        let result = bit_reader.read_signed_exp_golomb().unwrap();
321        assert_eq!(result, 2);
322        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
323
324        let result = bit_reader.read_signed_exp_golomb().unwrap();
325        assert_eq!(result, -2);
326        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
327
328        let result = bit_reader.read_signed_exp_golomb().unwrap();
329        assert_eq!(result, 3);
330        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
331
332        let result = bit_reader.read_signed_exp_golomb().unwrap();
333        assert_eq!(result, -3);
334        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
335
336        let result = bit_reader.read_signed_exp_golomb().unwrap();
337        assert_eq!(result, i64::MAX);
338        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 154);
339    }
340}