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#![deny(missing_docs)]
43#![deny(unsafe_code)]
44#![deny(unreachable_pub)]
45
46use std::io;
47
48use scuffle_bytes_util::{BitReader, BitWriter};
49
50/// Extension trait for reading Exp-Golomb encoded numbers from a bit reader
51///
52/// See: <https://en.wikipedia.org/wiki/Exponential-Golomb_coding>
53///
54/// - [`BitReader`]
55pub trait BitReaderExpGolombExt {
56    /// Reads an Exp-Golomb encoded number
57    fn read_exp_golomb(&mut self) -> io::Result<u64>;
58
59    /// Reads a signed Exp-Golomb encoded number
60    fn read_signed_exp_golomb(&mut self) -> io::Result<i64> {
61        let exp_glob = self.read_exp_golomb()?;
62
63        if exp_glob % 2 == 0 {
64            Ok(-((exp_glob / 2) as i64))
65        } else {
66            Ok((exp_glob / 2) as i64 + 1)
67        }
68    }
69}
70
71impl<R: io::Read> BitReaderExpGolombExt for BitReader<R> {
72    fn read_exp_golomb(&mut self) -> io::Result<u64> {
73        let mut leading_zeros = 0;
74        while !self.read_bit()? {
75            leading_zeros += 1;
76        }
77
78        let mut result = 1;
79        for _ in 0..leading_zeros {
80            result <<= 1;
81            result |= self.read_bit()? as u64;
82        }
83
84        Ok(result - 1)
85    }
86}
87
88/// Extension trait for writing Exp-Golomb encoded numbers to a bit writer
89///
90/// See: <https://en.wikipedia.org/wiki/Exponential-Golomb_coding>
91///
92/// - [`BitWriter`]
93pub trait BitWriterExpGolombExt {
94    /// Writes an Exp-Golomb encoded number
95    fn write_exp_golomb(&mut self, input: u64) -> io::Result<()>;
96
97    /// Writes a signed Exp-Golomb encoded number
98    fn write_signed_exp_golomb(&mut self, number: i64) -> io::Result<()> {
99        let number = if number <= 0 {
100            -number as u64 * 2
101        } else {
102            number as u64 * 2 - 1
103        };
104
105        self.write_exp_golomb(number)
106    }
107}
108
109impl<W: io::Write> BitWriterExpGolombExt for BitWriter<W> {
110    fn write_exp_golomb(&mut self, input: u64) -> io::Result<()> {
111        let mut number = input + 1;
112        let mut leading_zeros = 0;
113        while number > 1 {
114            number >>= 1;
115            leading_zeros += 1;
116        }
117
118        for _ in 0..leading_zeros {
119            self.write_bit(false)?;
120        }
121
122        self.write_bits(input + 1, leading_zeros + 1)?;
123
124        Ok(())
125    }
126}
127
128/// Returns the number of bits that a signed Exp-Golomb encoded number would take up.
129///
130/// See: <https://en.wikipedia.org/wiki/Exponential-Golomb_coding>
131pub fn size_of_signed_exp_golomb(number: i64) -> u64 {
132    let number = if number <= 0 {
133        -number as u64 * 2
134    } else {
135        number as u64 * 2 - 1
136    };
137
138    size_of_exp_golomb(number)
139}
140
141/// Returns the number of bits that an Exp-Golomb encoded number would take up.
142///
143/// See: <https://en.wikipedia.org/wiki/Exponential-Golomb_coding>
144pub fn size_of_exp_golomb(number: u64) -> u64 {
145    let mut number = number + 1;
146    let mut leading_zeros = 0;
147    while number > 1 {
148        number >>= 1;
149        leading_zeros += 1;
150    }
151
152    leading_zeros * 2 + 1
153}
154
155#[cfg(test)]
156#[cfg_attr(all(test, coverage_nightly), coverage(off))]
157mod tests {
158    use bytes::Buf;
159    use scuffle_bytes_util::{BitReader, BitWriter};
160
161    use crate::{BitReaderExpGolombExt, BitWriterExpGolombExt, size_of_exp_golomb, size_of_signed_exp_golomb};
162
163    pub(crate) fn get_remaining_bits(reader: &BitReader<std::io::Cursor<Vec<u8>>>) -> usize {
164        let remaining = reader.get_ref().remaining();
165
166        if reader.is_aligned() {
167            remaining * 8
168        } else {
169            remaining * 8 + (8 - reader.bit_pos() as usize)
170        }
171    }
172
173    #[test]
174    fn test_exp_glob_decode() {
175        let mut bit_writer = BitWriter::<Vec<u8>>::default();
176
177        bit_writer.write_bits(0b1, 1).unwrap(); // 0
178        bit_writer.write_bits(0b010, 3).unwrap(); // 1
179        bit_writer.write_bits(0b011, 3).unwrap(); // 2
180        bit_writer.write_bits(0b00100, 5).unwrap(); // 3
181        bit_writer.write_bits(0b00101, 5).unwrap(); // 4
182        bit_writer.write_bits(0b00110, 5).unwrap(); // 5
183        bit_writer.write_bits(0b00111, 5).unwrap(); // 6
184
185        let data = bit_writer.finish().unwrap();
186
187        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
188
189        let remaining_bits = get_remaining_bits(&bit_reader);
190
191        let result = bit_reader.read_exp_golomb().unwrap();
192        assert_eq!(result, 0);
193        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
194
195        let result = bit_reader.read_exp_golomb().unwrap();
196        assert_eq!(result, 1);
197        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
198
199        let result = bit_reader.read_exp_golomb().unwrap();
200        assert_eq!(result, 2);
201        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
202
203        let result = bit_reader.read_exp_golomb().unwrap();
204        assert_eq!(result, 3);
205        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
206
207        let result = bit_reader.read_exp_golomb().unwrap();
208        assert_eq!(result, 4);
209        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
210
211        let result = bit_reader.read_exp_golomb().unwrap();
212        assert_eq!(result, 5);
213        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
214
215        let result = bit_reader.read_exp_golomb().unwrap();
216        assert_eq!(result, 6);
217        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
218    }
219
220    #[test]
221    fn test_signed_exp_glob_decode() {
222        let mut bit_writer = BitWriter::<Vec<u8>>::default();
223
224        bit_writer.write_bits(0b1, 1).unwrap(); // 0
225        bit_writer.write_bits(0b010, 3).unwrap(); // 1
226        bit_writer.write_bits(0b011, 3).unwrap(); // -1
227        bit_writer.write_bits(0b00100, 5).unwrap(); // 2
228        bit_writer.write_bits(0b00101, 5).unwrap(); // -2
229        bit_writer.write_bits(0b00110, 5).unwrap(); // 3
230        bit_writer.write_bits(0b00111, 5).unwrap(); // -3
231
232        let data = bit_writer.finish().unwrap();
233
234        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
235
236        let remaining_bits = get_remaining_bits(&bit_reader);
237
238        let result = bit_reader.read_signed_exp_golomb().unwrap();
239        assert_eq!(result, 0);
240        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
241
242        let result = bit_reader.read_signed_exp_golomb().unwrap();
243        assert_eq!(result, 1);
244        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
245
246        let result = bit_reader.read_signed_exp_golomb().unwrap();
247        assert_eq!(result, -1);
248        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
249
250        let result = bit_reader.read_signed_exp_golomb().unwrap();
251        assert_eq!(result, 2);
252        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
253
254        let result = bit_reader.read_signed_exp_golomb().unwrap();
255        assert_eq!(result, -2);
256        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
257
258        let result = bit_reader.read_signed_exp_golomb().unwrap();
259        assert_eq!(result, 3);
260        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
261
262        let result = bit_reader.read_signed_exp_golomb().unwrap();
263        assert_eq!(result, -3);
264        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
265    }
266
267    #[test]
268    fn test_exp_glob_encode() {
269        let mut bit_writer = BitWriter::<Vec<u8>>::default();
270
271        bit_writer.write_exp_golomb(0).unwrap();
272        bit_writer.write_exp_golomb(1).unwrap();
273        bit_writer.write_exp_golomb(2).unwrap();
274        bit_writer.write_exp_golomb(3).unwrap();
275        bit_writer.write_exp_golomb(4).unwrap();
276        bit_writer.write_exp_golomb(5).unwrap();
277        bit_writer.write_exp_golomb(6).unwrap();
278        bit_writer.write_exp_golomb(u64::MAX - 1).unwrap();
279
280        let data = bit_writer.finish().unwrap();
281
282        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
283
284        let remaining_bits = get_remaining_bits(&bit_reader);
285
286        let result = bit_reader.read_exp_golomb().unwrap();
287        assert_eq!(result, 0);
288        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
289
290        let result = bit_reader.read_exp_golomb().unwrap();
291        assert_eq!(result, 1);
292        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
293
294        let result = bit_reader.read_exp_golomb().unwrap();
295        assert_eq!(result, 2);
296        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
297
298        let result = bit_reader.read_exp_golomb().unwrap();
299        assert_eq!(result, 3);
300        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
301
302        let result = bit_reader.read_exp_golomb().unwrap();
303        assert_eq!(result, 4);
304        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
305
306        let result = bit_reader.read_exp_golomb().unwrap();
307        assert_eq!(result, 5);
308        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
309
310        let result = bit_reader.read_exp_golomb().unwrap();
311        assert_eq!(result, 6);
312        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
313
314        let result = bit_reader.read_exp_golomb().unwrap();
315        assert_eq!(result, u64::MAX - 1);
316        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 154);
317    }
318
319    #[test]
320    fn test_signed_exp_glob_encode() {
321        let mut bit_writer = BitWriter::<Vec<u8>>::default();
322
323        bit_writer.write_signed_exp_golomb(0).unwrap();
324        bit_writer.write_signed_exp_golomb(1).unwrap();
325        bit_writer.write_signed_exp_golomb(-1).unwrap();
326        bit_writer.write_signed_exp_golomb(2).unwrap();
327        bit_writer.write_signed_exp_golomb(-2).unwrap();
328        bit_writer.write_signed_exp_golomb(3).unwrap();
329        bit_writer.write_signed_exp_golomb(-3).unwrap();
330        bit_writer.write_signed_exp_golomb(i64::MAX).unwrap();
331
332        let data = bit_writer.finish().unwrap();
333
334        let mut bit_reader = BitReader::new(std::io::Cursor::new(data));
335
336        let remaining_bits = get_remaining_bits(&bit_reader);
337
338        let result = bit_reader.read_signed_exp_golomb().unwrap();
339        assert_eq!(result, 0);
340        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 1);
341
342        let result = bit_reader.read_signed_exp_golomb().unwrap();
343        assert_eq!(result, 1);
344        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 4);
345
346        let result = bit_reader.read_signed_exp_golomb().unwrap();
347        assert_eq!(result, -1);
348        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 7);
349
350        let result = bit_reader.read_signed_exp_golomb().unwrap();
351        assert_eq!(result, 2);
352        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 12);
353
354        let result = bit_reader.read_signed_exp_golomb().unwrap();
355        assert_eq!(result, -2);
356        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 17);
357
358        let result = bit_reader.read_signed_exp_golomb().unwrap();
359        assert_eq!(result, 3);
360        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 22);
361
362        let result = bit_reader.read_signed_exp_golomb().unwrap();
363        assert_eq!(result, -3);
364        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 27);
365
366        let result = bit_reader.read_signed_exp_golomb().unwrap();
367        assert_eq!(result, i64::MAX);
368        assert_eq!(get_remaining_bits(&bit_reader), remaining_bits - 154);
369    }
370
371    #[test]
372    fn test_expg_sizes() {
373        assert_eq!(1, size_of_exp_golomb(0)); // 0b1
374        assert_eq!(3, size_of_exp_golomb(1)); // 0b010
375        assert_eq!(3, size_of_exp_golomb(2)); // 0b011
376        assert_eq!(5, size_of_exp_golomb(3)); // 0b00100
377        assert_eq!(5, size_of_exp_golomb(4)); // 0b00101
378        assert_eq!(5, size_of_exp_golomb(5)); // 0b00110
379        assert_eq!(5, size_of_exp_golomb(6)); // 0b00111
380
381        assert_eq!(1, size_of_signed_exp_golomb(0)); // 0b1
382        assert_eq!(3, size_of_signed_exp_golomb(1)); // 0b010
383        assert_eq!(3, size_of_signed_exp_golomb(-1)); // 0b011
384        assert_eq!(5, size_of_signed_exp_golomb(2)); // 0b00100
385        assert_eq!(5, size_of_signed_exp_golomb(-2)); // 0b00101
386        assert_eq!(5, size_of_signed_exp_golomb(3)); // 0b00110
387        assert_eq!(5, size_of_signed_exp_golomb(-3)); // 0b00111
388    }
389}