exp_golomb/
encoder.rs

1/// An Exponential-Golomb writer.
2pub struct ExpGolombEncoder<'a> {
3    bit_buf: BitBuffer<'a>,
4}
5
6impl<'a> ExpGolombEncoder<'a> {
7    /// Create a new `ExpGolombEncoder`.
8    ///
9    /// `start` denotes the starting position in the first byte of `buf` and goes from 0 (first) to
10    ///  7 (last). This function returns `None` if the buffer is empty or if `start` is  not within
11    /// \[0, 7\].
12    /// 
13    /// # Examples
14    ///
15    /// ```
16    /// # use exp_golomb::ExpGolombEncoder;
17    /// let mut buf = [0u8; 1];
18    /// // Write starting at the second bit
19    /// let mut writer = ExpGolombEncoder::new(&mut buf, 1).unwrap();
20    /// writer.put_unsigned(2).unwrap();
21    /// writer.close();
22    /// assert_eq!(buf[0], 0b00110000);
23    /// ```
24    #[inline]
25    #[must_use]
26    pub fn new(buf: &'a mut [u8], start: u32) -> Option<ExpGolombEncoder<'a>> {
27        if buf.is_empty() || start > 7 {
28            return None;
29        }
30        Some(ExpGolombEncoder {
31            bit_buf: BitBuffer::new(buf, start),
32        })
33    }
34
35    /// Encode a `u64` into the buffer. Returns `None` if the buffer is full.
36    ///
37    /// # Examples
38    ///
39    /// ```
40    /// # use exp_golomb::ExpGolombEncoder;
41    /// let mut buf = [0u8; 6];
42    /// let mut writer = ExpGolombEncoder::new(&mut buf, 0).unwrap();
43    /// for i in 0..=8 {
44    ///     writer.put_unsigned(i).unwrap();
45    /// }
46    /// writer.close();
47    /// 
48    /// assert_eq!(
49    ///     buf,
50    ///     [0b10100110, 0b01000010, 0b10011000, 0b11100010, 0b00000100, 0b10000000]
51    /// );
52    /// ```
53    /// 
54    /// This function guards against out of bounds indexing by returning `None`:
55    /// 
56    /// ```
57    /// # use exp_golomb::ExpGolombEncoder;
58    /// let mut buf = [0u8; 1];
59    /// let mut writer = ExpGolombEncoder::new(&mut buf, 0).unwrap();
60    /// assert!(writer.put_unsigned(1).is_some());
61    /// assert!(writer.put_unsigned(1).is_some());
62    /// assert!(writer.put_unsigned(1).is_none());
63    /// ```
64    #[inline]
65    #[must_use]
66    pub fn put_unsigned(&mut self, value: u64) -> Option<()> {
67        let xp1 = value.wrapping_add(1);
68
69        let bytes = xp1.to_be_bytes();
70        let lz = xp1.leading_zeros();
71        let start = (lz / 8) as usize;
72        let bit_start = lz - (lz / 8 * 8);
73
74        let num_zeros = 64 - lz - 1;
75        self.bit_buf.put_zeros(num_zeros);
76
77        self.bit_buf.put_bytes(&bytes[start..], bit_start)
78    }
79
80    /// Write a single bit to the buffer. Returns `None` if the buffer is full.
81    /// 
82    /// # Examples
83    ///
84    /// ```
85    /// # use exp_golomb::ExpGolombEncoder;
86    /// let mut buf = [0u8; 1];
87    /// let mut writer = ExpGolombEncoder::new(&mut buf, 6).unwrap();
88    /// writer.put_bit(true).unwrap();
89    /// writer.put_bit(false).unwrap();
90    /// assert!(writer.put_bit(true).is_none());
91    /// assert!(writer.put_bit(true).is_none());
92    /// writer.close();
93    /// assert_eq!(buf[0], 0b00000010);
94    /// ```
95    #[inline]
96    #[must_use]
97    pub fn put_bit(&mut self, value: bool) -> Option<()> {
98        self.bit_buf.put_bit(value)
99    }
100
101    /// Consumes the `ExpGolombEncoder`, returning the bit position one past the last written bit.
102    /// 
103    /// # Examples
104    ///
105    /// ```
106    /// # use exp_golomb::ExpGolombEncoder;
107    /// let mut buf = [0u8; 1];
108    /// let mut writer = ExpGolombEncoder::new(&mut buf, 2).unwrap();
109    /// writer.put_unsigned(0).unwrap();
110    /// assert_eq!(writer.close(), (0, 3));
111    /// ```
112    #[inline]
113    pub fn close(self) -> (usize, u32) {
114        (self.bit_buf.index, self.bit_buf.bit_pos)
115    }
116}
117
118struct BitBuffer<'a> {
119    buf: &'a mut [u8],
120    index: usize,
121    bit_pos: u32,
122}
123
124impl<'a> BitBuffer<'a> {
125    #[inline]
126    fn new(buf: &'a mut [u8], bit_pos: u32) -> BitBuffer<'a> {
127        BitBuffer {
128            buf,
129            index: 0,
130            bit_pos,
131        }
132    }
133
134    #[inline]
135    fn put_bit(&mut self, value: bool) -> Option<()> {
136        *self.buf.get_mut(self.index)? |= (value as u8) << (7 - self.bit_pos);
137        self.bit_pos += 1;
138        if self.bit_pos >= 8 {
139            self.bit_pos -= 8;
140            self.index += 1;
141        }
142        Some(())
143    }
144
145    #[inline]
146    fn put_zeros(&mut self, num_zeros: u32) -> Option<()> {
147        // TODO: Suboptimal
148        for _ in 0..num_zeros {
149            self.put_bit(false)?;
150        }
151        Some(())
152    }
153
154    #[inline]
155    #[must_use]
156    fn put_bytes(&mut self, bytes: &[u8], mut start_pos: u32) -> Option<()> {
157        for &byte in bytes {
158            while start_pos < 8 {
159                let data = ((byte as u32) << start_pos) >> self.bit_pos;
160                *self.buf.get_mut(self.index)? |= data as u8;
161
162                let shift_amount = 8 - u32::max(self.bit_pos, start_pos);
163                self.bit_pos += shift_amount;
164                if self.bit_pos >= 8 {
165                    self.bit_pos -= 8;
166                    self.index += 1;
167                }
168
169                start_pos += shift_amount;
170            }
171            start_pos -= 8;
172        }
173        Some(())
174    }
175}
exp_golomb/encoder.rs

exp_golomb/
encoder.rs
