Skip to main content

bincode_next/enc/
bit_writer.rs

1use crate::enc::write::Writer;
2use crate::error::EncodeError;
3
4/// A wrapper around a `Writer` that allows writing data bit-by-bit.
5pub struct BitWriter<'a, W: Writer> {
6    writer: &'a mut W,
7    current_byte: u8,
8    bit_count: u8,
9}
10
11impl<'a, W: Writer> BitWriter<'a, W> {
12    /// Creates a new `BitWriter`.
13    #[inline(always)]
14    pub const fn new(writer: &'a mut W) -> Self {
15        Self {
16            writer,
17            current_byte: 0,
18            bit_count: 0,
19        }
20    }
21
22    /// Creates a new `BitWriter` with an existing state.
23    #[inline(always)]
24    pub const fn from_state(
25        writer: &'a mut W,
26        current_byte: u8,
27        bit_count: u8,
28    ) -> Self {
29        Self {
30            writer,
31            current_byte,
32            bit_count,
33        }
34    }
35
36    /// Returns the current state of the bit writer (`current_byte`, `bit_count`).
37    #[inline(always)]
38    #[must_use]
39    pub const fn get_state(&self) -> (u8, u8) {
40        (self.current_byte, self.bit_count)
41    }
42
43    /// Writes `num_bits` from `val` into the stream, using LSB-first bit ordering.
44    ///
45    /// # Errors
46    ///
47    /// Returns `EncodeError` if the underlying writer fails to write the data.
48    #[inline(always)]
49    pub fn write_bits_lsb(
50        &mut self,
51        mut val: u64,
52        mut num_bits: u8,
53    ) -> Result<(), EncodeError> {
54        while num_bits > 0 {
55            let space_in_byte = 8 - self.bit_count;
56            let bits_to_write = num_bits.min(space_in_byte);
57
58            let mask = (1u64 << bits_to_write) - 1;
59            let chunk = (val & mask) as u8;
60
61            self.current_byte |= chunk << self.bit_count;
62            self.bit_count += bits_to_write;
63
64            val >>= bits_to_write;
65            num_bits -= bits_to_write;
66
67            if self.bit_count == 8 {
68                self.flush()?;
69            }
70        }
71        Ok(())
72    }
73
74    /// Writes `num_bits` from `val` into the stream, using MSB-first bit ordering.
75    ///
76    /// # Errors
77    ///
78    /// Returns `EncodeError` if the underlying writer fails to write the data.
79    #[inline(always)]
80    pub fn write_bits_msb(
81        &mut self,
82        val: u64,
83        mut num_bits: u8,
84    ) -> Result<(), EncodeError> {
85        while num_bits > 0 {
86            let space_in_byte = 8 - self.bit_count;
87            let bits_to_write = num_bits.min(space_in_byte);
88
89            // Extract the top `bits_to_write` from the `num_bits` we care about in `val`
90            let shift_down = num_bits - bits_to_write;
91            let mask = (1u64 << bits_to_write) - 1;
92            let chunk = ((val >> shift_down) & mask) as u8;
93
94            let shift_up = 8 - self.bit_count - bits_to_write;
95            self.current_byte |= chunk << shift_up;
96            self.bit_count += bits_to_write;
97
98            num_bits -= bits_to_write;
99
100            if self.bit_count == 8 {
101                self.flush()?;
102            }
103        }
104        Ok(())
105    }
106
107    /// Writes `num_bits` from `val` into the stream, using the bit ordering from the configuration.
108    ///
109    /// # Errors
110    ///
111    /// Returns `EncodeError` if the underlying writer fails to write the data.
112    #[inline(always)]
113    pub fn write_bits<C: crate::config::Config>(
114        &mut self,
115        val: u64,
116        num_bits: u8,
117        config: &C,
118    ) -> Result<(), EncodeError> {
119        use crate::config::BitOrdering;
120        match config.bit_ordering() {
121            | BitOrdering::Lsb => self.write_bits_lsb(val, num_bits),
122            | BitOrdering::Msb => self.write_bits_msb(val, num_bits),
123        }
124    }
125
126    /// Flushes any remaining bits to the underlying writer, padding with 0s.
127    ///
128    /// # Errors
129    ///
130    /// Returns `EncodeError` if the underlying writer fails to write the data.
131    #[inline(always)]
132    pub fn flush(&mut self) -> Result<(), EncodeError> {
133        if self.bit_count > 0 {
134            self.writer.write(&[self.current_byte])?;
135            self.current_byte = 0;
136            self.bit_count = 0;
137        }
138        Ok(())
139    }
140}