Skip to main content

vexil_runtime/
bit_writer.rs

1pub struct BitWriter {
2    buf: Vec<u8>,
3    current_byte: u8,
4    bit_offset: u8,
5}
6
7impl BitWriter {
8    pub fn new() -> Self {
9        Self {
10            buf: Vec::new(),
11            current_byte: 0,
12            bit_offset: 0,
13        }
14    }
15
16    /// Internal: align to a byte boundary without the "empty = zero byte" rule.
17    /// Used before multi-byte writes to ensure alignment.
18    fn align(&mut self) {
19        if self.bit_offset > 0 {
20            self.buf.push(self.current_byte);
21            self.current_byte = 0;
22            self.bit_offset = 0;
23        }
24    }
25
26    /// Write `count` bits from `value`, LSB first.
27    pub fn write_bits(&mut self, value: u64, count: u8) {
28        let mut v = value;
29        for _ in 0..count {
30            let bit = (v & 1) as u8;
31            self.current_byte |= bit << self.bit_offset;
32            self.bit_offset += 1;
33            if self.bit_offset == 8 {
34                self.buf.push(self.current_byte);
35                self.current_byte = 0;
36                self.bit_offset = 0;
37            }
38            v >>= 1;
39        }
40    }
41
42    /// Write a single boolean as 1 bit.
43    pub fn write_bool(&mut self, v: bool) {
44        self.write_bits(v as u64, 1);
45    }
46
47    /// Flush any partial byte to the buffer.
48    ///
49    /// Special case per spec §4.1: if nothing has been written at all
50    /// (bit_offset == 0 AND buf is empty), push a zero byte anyway.
51    /// If bit_offset == 0 and buf is non-empty, this is a no-op.
52    pub fn flush_to_byte_boundary(&mut self) {
53        if self.bit_offset == 0 {
54            if self.buf.is_empty() {
55                self.buf.push(0x00);
56            }
57            // else: already aligned and something was written — no-op
58        } else {
59            self.buf.push(self.current_byte);
60            self.current_byte = 0;
61            self.bit_offset = 0;
62        }
63    }
64
65    pub fn write_u8(&mut self, v: u8) {
66        self.align();
67        self.buf.push(v);
68    }
69
70    pub fn write_u16(&mut self, v: u16) {
71        self.align();
72        self.buf.extend_from_slice(&v.to_le_bytes());
73    }
74
75    pub fn write_u32(&mut self, v: u32) {
76        self.align();
77        self.buf.extend_from_slice(&v.to_le_bytes());
78    }
79
80    pub fn write_u64(&mut self, v: u64) {
81        self.align();
82        self.buf.extend_from_slice(&v.to_le_bytes());
83    }
84
85    pub fn write_i8(&mut self, v: i8) {
86        self.align();
87        self.buf.extend_from_slice(&v.to_le_bytes());
88    }
89
90    pub fn write_i16(&mut self, v: i16) {
91        self.align();
92        self.buf.extend_from_slice(&v.to_le_bytes());
93    }
94
95    pub fn write_i32(&mut self, v: i32) {
96        self.align();
97        self.buf.extend_from_slice(&v.to_le_bytes());
98    }
99
100    pub fn write_i64(&mut self, v: i64) {
101        self.align();
102        self.buf.extend_from_slice(&v.to_le_bytes());
103    }
104
105    /// Write an f32, canonicalizing NaN to 0x7FC00000.
106    pub fn write_f32(&mut self, v: f32) {
107        self.align();
108        let bits: u32 = if v.is_nan() {
109            0x7FC00000u32
110        } else {
111            v.to_bits()
112        };
113        self.buf.extend_from_slice(&bits.to_le_bytes());
114    }
115
116    /// Write an f64, canonicalizing NaN to 0x7FF8000000000000.
117    pub fn write_f64(&mut self, v: f64) {
118        self.align();
119        let bits: u64 = if v.is_nan() {
120            0x7FF8000000000000u64
121        } else {
122            v.to_bits()
123        };
124        self.buf.extend_from_slice(&bits.to_le_bytes());
125    }
126
127    /// Write a LEB128-encoded unsigned integer.
128    pub fn write_leb128(&mut self, v: u64) {
129        self.align();
130        crate::leb128::encode(&mut self.buf, v);
131    }
132
133    /// Write a ZigZag + LEB128 encoded signed integer.
134    pub fn write_zigzag(&mut self, v: i64, type_bits: u8) {
135        let encoded = crate::zigzag::zigzag_encode(v, type_bits);
136        self.write_leb128(encoded);
137    }
138
139    /// Write a UTF-8 string with a LEB128 length prefix.
140    pub fn write_string(&mut self, s: &str) {
141        self.align();
142        crate::leb128::encode(&mut self.buf, s.len() as u64);
143        self.buf.extend_from_slice(s.as_bytes());
144    }
145
146    /// Write a byte slice with a LEB128 length prefix.
147    pub fn write_bytes(&mut self, data: &[u8]) {
148        self.align();
149        crate::leb128::encode(&mut self.buf, data.len() as u64);
150        self.buf.extend_from_slice(data);
151    }
152
153    /// Write raw bytes with no length prefix.
154    pub fn write_raw_bytes(&mut self, data: &[u8]) {
155        self.align();
156        self.buf.extend_from_slice(data);
157    }
158
159    /// Flush any partial byte and return the finished buffer.
160    pub fn finish(mut self) -> Vec<u8> {
161        self.flush_to_byte_boundary();
162        self.buf
163    }
164}
165
166impl Default for BitWriter {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175
176    #[test]
177    fn write_single_bit_true() {
178        let mut w = BitWriter::new();
179        w.write_bool(true);
180        assert_eq!(w.finish(), [0x01]);
181    }
182
183    #[test]
184    fn write_single_bit_false() {
185        let mut w = BitWriter::new();
186        w.write_bool(false);
187        assert_eq!(w.finish(), [0x00]);
188    }
189
190    #[test]
191    fn write_bits_lsb_first() {
192        let mut w = BitWriter::new();
193        w.write_bits(5, 3); // 101
194        w.write_bits(19, 5); // 10011
195                             // LSB-first: byte = 10011_101 = 0x9D
196        assert_eq!(w.finish(), [0x9D]);
197    }
198
199    #[test]
200    fn write_bits_cross_byte_boundary() {
201        let mut w = BitWriter::new();
202        w.write_bits(5, 3);
203        w.write_bits(19, 5);
204        w.write_bits(42, 6); // 101010
205                             // Byte 0: 0x9D, Byte 1: 00_101010 = 0x2A
206        assert_eq!(w.finish(), [0x9D, 0x2A]);
207    }
208
209    #[test]
210    fn flush_to_byte_boundary_pads_zeros() {
211        let mut w = BitWriter::new();
212        w.write_bits(0b101, 3);
213        w.flush_to_byte_boundary();
214        w.write_bits(0xFF, 8);
215        assert_eq!(w.finish(), [0x05, 0xFF]);
216    }
217
218    #[test]
219    fn write_u8_flushes_first() {
220        let mut w = BitWriter::new();
221        w.write_bool(true);
222        w.write_u8(0xAB);
223        assert_eq!(w.finish(), [0x01, 0xAB]);
224    }
225
226    #[test]
227    fn write_u16_le() {
228        let mut w = BitWriter::new();
229        w.write_u16(0x0102);
230        assert_eq!(w.finish(), [0x02, 0x01]);
231    }
232
233    #[test]
234    fn write_u32_le() {
235        let mut w = BitWriter::new();
236        w.write_u32(0x01020304);
237        assert_eq!(w.finish(), [0x04, 0x03, 0x02, 0x01]);
238    }
239
240    #[test]
241    fn write_i16_negative() {
242        let mut w = BitWriter::new();
243        w.write_i16(-1);
244        assert_eq!(w.finish(), [0xFF, 0xFF]);
245    }
246
247    #[test]
248    fn write_f32_nan_canonicalized() {
249        let mut w = BitWriter::new();
250        w.write_f32(f32::NAN);
251        assert_eq!(w.finish(), [0x00, 0x00, 0xC0, 0x7F]);
252    }
253
254    #[test]
255    fn write_f64_nan_canonicalized() {
256        let mut w = BitWriter::new();
257        w.write_f64(f64::NAN);
258        assert_eq!(w.finish(), 0x7FF8000000000000u64.to_le_bytes());
259    }
260
261    #[test]
262    fn write_f32_negative_zero_preserved() {
263        let mut w = BitWriter::new();
264        w.write_f32(-0.0f32);
265        let buf = w.finish();
266        assert_eq!(buf, (-0.0f32).to_le_bytes());
267        assert_ne!(buf, 0.0f32.to_le_bytes());
268    }
269
270    #[test]
271    fn write_leb128_test() {
272        let mut w = BitWriter::new();
273        w.write_leb128(300);
274        assert_eq!(w.finish(), [0xAC, 0x02]);
275    }
276
277    #[test]
278    fn write_zigzag_neg1() {
279        let mut w = BitWriter::new();
280        w.write_zigzag(-1, 64);
281        assert_eq!(w.finish(), [0x01]);
282    }
283
284    #[test]
285    fn write_string_test() {
286        let mut w = BitWriter::new();
287        w.write_string("hi");
288        assert_eq!(w.finish(), [0x02, 0x68, 0x69]);
289    }
290
291    #[test]
292    fn write_bytes_test() {
293        let mut w = BitWriter::new();
294        w.write_bytes(&[0xDE, 0xAD]);
295        assert_eq!(w.finish(), [0x02, 0xDE, 0xAD]);
296    }
297
298    #[test]
299    fn write_raw_bytes_test() {
300        let mut w = BitWriter::new();
301        w.write_raw_bytes(&[0xCA, 0xFE]);
302        assert_eq!(w.finish(), [0xCA, 0xFE]);
303    }
304
305    #[test]
306    fn empty_flush_produces_zero_byte() {
307        let mut w = BitWriter::new();
308        w.flush_to_byte_boundary();
309        assert_eq!(w.finish(), [0x00]);
310    }
311}