naia_serde/
bit_writer.rs

1use crate::{
2    constants::{MTU_SIZE_BITS, MTU_SIZE_BYTES},
3    BitCounter, OutgoingPacket, OwnedBitReader,
4};
5
6// BitWrite
7pub trait BitWrite {
8    fn write_bit(&mut self, bit: bool);
9    fn write_byte(&mut self, byte: u8);
10
11    fn is_counter(&self) -> bool;
12    fn count_bits(&mut self, bits: u32);
13}
14
15// BitWriter
16//
17// Internals: bits accumulate LSB-first into a u32 scratch register and flush
18// as a little-endian u32 word every 32 bits. Using u32 (not u64) keeps all
19// operations native on wasm32 targets where 64-bit arithmetic is emulated.
20// The approach eliminates the per-byte reverse_bits call of the old u8 design.
21// (Inspired by Gaffer on Games, "Reading and Writing Packets", 2015.)
22pub struct BitWriter {
23    scratch: u32,
24    scratch_bits: u32,
25    buffer: [u8; MTU_SIZE_BYTES],
26    byte_count: usize,
27    current_bits: u32,
28    max_bits: u32,
29}
30
31impl BitWriter {
32    #[allow(clippy::new_without_default)]
33    pub fn new() -> Self {
34        Self {
35            scratch: 0,
36            scratch_bits: 0,
37            buffer: [0; MTU_SIZE_BYTES],
38            byte_count: 0,
39            current_bits: 0,
40            max_bits: MTU_SIZE_BITS,
41        }
42    }
43
44    pub fn with_capacity(bit_capacity: u32) -> Self {
45        Self {
46            scratch: 0,
47            scratch_bits: 0,
48            buffer: [0; MTU_SIZE_BYTES],
49            byte_count: 0,
50            current_bits: 0,
51            max_bits: bit_capacity,
52        }
53    }
54
55    pub fn with_max_capacity() -> Self {
56        Self::with_capacity(u32::MAX)
57    }
58
59    fn flush_word(&mut self) {
60        self.buffer[self.byte_count..self.byte_count + 4]
61            .copy_from_slice(&self.scratch.to_le_bytes());
62        self.byte_count += 4;
63        self.scratch = 0;
64        self.scratch_bits = 0;
65    }
66
67    fn finalize(&mut self) {
68        if self.scratch_bits > 0 {
69            let remaining_bytes = (self.scratch_bits as usize).div_ceil(8);
70            let word = self.scratch.to_le_bytes();
71            self.buffer[self.byte_count..self.byte_count + remaining_bytes]
72                .copy_from_slice(&word[..remaining_bytes]);
73            self.byte_count += remaining_bytes;
74        }
75        self.max_bits = 0;
76    }
77
78    pub fn to_packet(mut self) -> OutgoingPacket {
79        self.finalize();
80        OutgoingPacket::new(self.byte_count, self.buffer)
81    }
82
83    pub fn to_owned_reader(mut self) -> OwnedBitReader {
84        self.finalize();
85        OwnedBitReader::new(&self.buffer[0..self.byte_count])
86    }
87
88    pub fn to_bytes(mut self) -> Box<[u8]> {
89        self.finalize();
90        Box::from(&self.buffer[0..self.byte_count])
91    }
92
93    pub fn counter(&self) -> BitCounter {
94        BitCounter::new(self.current_bits, self.current_bits, self.max_bits)
95    }
96
97    pub fn reserve_bits(&mut self, bits: u32) {
98        self.max_bits -= bits;
99    }
100
101    pub fn release_bits(&mut self, bits: u32) {
102        self.max_bits += bits;
103    }
104
105    pub fn bits_free(&self) -> u32 {
106        self.max_bits - self.current_bits
107    }
108}
109
110impl BitWrite for BitWriter {
111    #[inline(always)]
112    fn write_bit(&mut self, bit: bool) {
113        if self.current_bits >= self.max_bits {
114            panic!("Write overflow!");
115        }
116        self.scratch |= (bit as u32) << self.scratch_bits;
117        self.scratch_bits += 1;
118        self.current_bits += 1;
119        if self.scratch_bits == 32 {
120            self.flush_word();
121        }
122    }
123
124    #[inline(always)]
125    fn write_byte(&mut self, byte: u8) {
126        if self.current_bits + 8 > self.max_bits {
127            panic!("Write overflow!");
128        }
129        self.current_bits += 8;
130        let available = 32 - self.scratch_bits;
131        if available >= 8 {
132            self.scratch |= (byte as u32) << self.scratch_bits;
133            self.scratch_bits += 8;
134            if self.scratch_bits == 32 {
135                self.flush_word();
136            }
137        } else {
138            // byte spans a 32-bit word boundary
139            let lo = (byte as u32) & ((1 << available) - 1);
140            self.scratch |= lo << self.scratch_bits;
141            self.flush_word();
142            self.scratch = (byte as u32) >> available;
143            self.scratch_bits = 8 - available;
144        }
145    }
146
147    fn count_bits(&mut self, _: u32) {
148        panic!("This method should not be called for BitWriter!");
149    }
150
151    fn is_counter(&self) -> bool {
152        false
153    }
154}
155
156mod tests {
157
158    // ─── word-boundary regression tests (targets for word-aligned optimization) ─
159
160    #[test]
161    fn read_write_33_bits() {
162        // 33 bits spans the 32-bit word boundary in the new implementation.
163        use crate::{bit_reader::BitReader, bit_writer::{BitWrite, BitWriter}};
164        let mut writer = BitWriter::with_max_capacity();
165        // write 33 known bits: alternating pattern
166        for i in 0..33usize {
167            writer.write_bit(i % 3 == 0);
168        }
169        let buffer = writer.to_bytes();
170        let mut reader = BitReader::new(&buffer);
171        for i in 0..33usize {
172            assert_eq!(reader.read_bit().unwrap(), i % 3 == 0, "bit {i} mismatch");
173        }
174    }
175
176    #[test]
177    fn read_write_64_bits_exact() {
178        // exactly 2 words — tests two full-word flushes
179        use crate::{bit_reader::BitReader, bit_writer::{BitWrite, BitWriter}};
180        let mut writer = BitWriter::with_max_capacity();
181        for i in 0..64usize {
182            writer.write_bit(i % 5 < 2);
183        }
184        let buffer = writer.to_bytes();
185        let mut reader = BitReader::new(&buffer);
186        for i in 0..64usize {
187            assert_eq!(reader.read_bit().unwrap(), i % 5 < 2, "bit {i} mismatch");
188        }
189    }
190
191    #[test]
192    fn read_write_5_bytes_via_write_byte_then_read_bit() {
193        // mix write_byte (8 aligned) with read_bit to verify no endian confusion
194        use crate::{bit_reader::BitReader, bit_writer::{BitWrite, BitWriter}};
195        let data: &[u8] = &[0b10110001, 0b01001110, 0b11010101, 0b00110011, 0b11111010];
196        let mut writer = BitWriter::with_max_capacity();
197        for &b in data {
198            writer.write_byte(b);
199        }
200        let buffer = writer.to_bytes();
201        let mut reader = BitReader::new(&buffer);
202        for &b in data {
203            for bit in 0..8usize {
204                let expected = (b >> bit) & 1 != 0;
205                assert_eq!(reader.read_bit().unwrap(), expected, "byte {b:#010b} bit {bit}");
206            }
207        }
208    }
209
210    // ─── existing bit/byte round-trip tests ────────────────────────────────────
211
212    #[test]
213    fn read_write_1_bit() {
214        use crate::{
215            bit_reader::BitReader,
216            bit_writer::{BitWrite, BitWriter},
217        };
218
219        let mut writer = BitWriter::new();
220
221        writer.write_bit(true);
222
223        let buffer = writer.to_bytes();
224
225        let mut reader = BitReader::new(&buffer);
226
227        assert!(reader.read_bit().unwrap());
228    }
229
230    #[test]
231    fn read_write_3_bits() {
232        use crate::{
233            bit_reader::BitReader,
234            bit_writer::{BitWrite, BitWriter},
235        };
236
237        let mut writer = BitWriter::new();
238
239        writer.write_bit(false);
240        writer.write_bit(true);
241        writer.write_bit(true);
242
243        let buffer = writer.to_bytes();
244
245        let mut reader = BitReader::new(&buffer);
246
247        assert!(!reader.read_bit().unwrap());
248        assert!(reader.read_bit().unwrap());
249        assert!(reader.read_bit().unwrap());
250    }
251
252    #[test]
253    fn read_write_8_bits() {
254        use crate::{
255            bit_reader::BitReader,
256            bit_writer::{BitWrite, BitWriter},
257        };
258
259        let mut writer = BitWriter::new();
260
261        writer.write_bit(false);
262        writer.write_bit(true);
263        writer.write_bit(false);
264        writer.write_bit(true);
265
266        writer.write_bit(true);
267        writer.write_bit(false);
268        writer.write_bit(false);
269        writer.write_bit(false);
270
271        let buffer = writer.to_bytes();
272
273        let mut reader = BitReader::new(&buffer);
274
275        assert!(!reader.read_bit().unwrap());
276        assert!(reader.read_bit().unwrap());
277        assert!(!reader.read_bit().unwrap());
278        assert!(reader.read_bit().unwrap());
279
280        assert!(reader.read_bit().unwrap());
281        assert!(!reader.read_bit().unwrap());
282        assert!(!reader.read_bit().unwrap());
283        assert!(!reader.read_bit().unwrap());
284    }
285
286    #[test]
287    fn read_write_13_bits() {
288        use crate::{
289            bit_reader::BitReader,
290            bit_writer::{BitWrite, BitWriter},
291        };
292
293        let mut writer = BitWriter::new();
294
295        writer.write_bit(false);
296        writer.write_bit(true);
297        writer.write_bit(false);
298        writer.write_bit(true);
299
300        writer.write_bit(true);
301        writer.write_bit(false);
302        writer.write_bit(false);
303        writer.write_bit(false);
304
305        writer.write_bit(true);
306        writer.write_bit(false);
307        writer.write_bit(true);
308        writer.write_bit(true);
309
310        writer.write_bit(true);
311
312        let buffer = writer.to_bytes();
313
314        let mut reader = BitReader::new(&buffer);
315
316        assert!(!reader.read_bit().unwrap());
317        assert!(reader.read_bit().unwrap());
318        assert!(!reader.read_bit().unwrap());
319        assert!(reader.read_bit().unwrap());
320
321        assert!(reader.read_bit().unwrap());
322        assert!(!reader.read_bit().unwrap());
323        assert!(!reader.read_bit().unwrap());
324        assert!(!reader.read_bit().unwrap());
325
326        assert!(reader.read_bit().unwrap());
327        assert!(!reader.read_bit().unwrap());
328        assert!(reader.read_bit().unwrap());
329        assert!(reader.read_bit().unwrap());
330
331        assert!(reader.read_bit().unwrap());
332    }
333
334    #[test]
335    fn read_write_16_bits() {
336        use crate::{
337            bit_reader::BitReader,
338            bit_writer::{BitWrite, BitWriter},
339        };
340
341        let mut writer = BitWriter::new();
342
343        writer.write_bit(false);
344        writer.write_bit(true);
345        writer.write_bit(false);
346        writer.write_bit(true);
347
348        writer.write_bit(true);
349        writer.write_bit(false);
350        writer.write_bit(false);
351        writer.write_bit(false);
352
353        writer.write_bit(true);
354        writer.write_bit(false);
355        writer.write_bit(true);
356        writer.write_bit(true);
357
358        writer.write_bit(true);
359        writer.write_bit(false);
360        writer.write_bit(true);
361        writer.write_bit(true);
362
363        let buffer = writer.to_bytes();
364
365        let mut reader = BitReader::new(&buffer);
366
367        assert!(!reader.read_bit().unwrap());
368        assert!(reader.read_bit().unwrap());
369        assert!(!reader.read_bit().unwrap());
370        assert!(reader.read_bit().unwrap());
371
372        assert!(reader.read_bit().unwrap());
373        assert!(!reader.read_bit().unwrap());
374        assert!(!reader.read_bit().unwrap());
375        assert!(!reader.read_bit().unwrap());
376
377        assert!(reader.read_bit().unwrap());
378        assert!(!reader.read_bit().unwrap());
379        assert!(reader.read_bit().unwrap());
380        assert!(reader.read_bit().unwrap());
381
382        assert!(reader.read_bit().unwrap());
383        assert!(!reader.read_bit().unwrap());
384        assert!(reader.read_bit().unwrap());
385        assert!(reader.read_bit().unwrap());
386    }
387
388    #[test]
389    fn read_write_1_byte() {
390        use crate::{
391            bit_reader::BitReader,
392            bit_writer::{BitWrite, BitWriter},
393        };
394
395        let mut writer = BitWriter::new();
396
397        writer.write_byte(123);
398
399        let buffer = writer.to_bytes();
400
401        let mut reader = BitReader::new(&buffer);
402
403        assert_eq!(123, reader.read_byte().unwrap());
404    }
405
406    #[test]
407    fn read_write_5_bytes() {
408        use crate::{
409            bit_reader::BitReader,
410            bit_writer::{BitWrite, BitWriter},
411        };
412
413        let mut writer = BitWriter::new();
414
415        writer.write_byte(48);
416        writer.write_byte(151);
417        writer.write_byte(62);
418        writer.write_byte(34);
419        writer.write_byte(2);
420
421        let buffer = writer.to_bytes();
422
423        let mut reader = BitReader::new(&buffer);
424
425        assert_eq!(48, reader.read_byte().unwrap());
426        assert_eq!(151, reader.read_byte().unwrap());
427        assert_eq!(62, reader.read_byte().unwrap());
428        assert_eq!(34, reader.read_byte().unwrap());
429        assert_eq!(2, reader.read_byte().unwrap());
430    }
431}
naia_serde/bit_writer.rs

naia_serde/
bit_writer.rs
