zrip_core/bitstream/
writer.rs

1#![forbid(unsafe_code)]
2
3#[cfg(feature = "alloc")]
4use alloc::vec::Vec;
5
6use super::primitives;
7
8pub struct BitWriter {
9    #[cfg(feature = "alloc")]
10    buf: Vec<u8>,
11    bits: u64,
12    bits_used: u8,
13    pos: usize,
14}
15
16#[cfg(feature = "alloc")]
17impl Default for BitWriter {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23#[cfg(feature = "alloc")]
24impl BitWriter {
25    pub fn new() -> Self {
26        Self {
27            buf: Vec::with_capacity(64),
28            bits: 0,
29            bits_used: 0,
30            pos: 0,
31        }
32    }
33
34    pub fn with_capacity(cap: usize) -> Self {
35        Self {
36            buf: Vec::with_capacity(cap + 8),
37            bits: 0,
38            bits_used: 0,
39            pos: 0,
40        }
41    }
42
43    pub fn from_vec(mut buf: Vec<u8>) -> Self {
44        buf.clear();
45        buf.reserve(8);
46        Self {
47            buf,
48            bits: 0,
49            bits_used: 0,
50            pos: 0,
51        }
52    }
53
54    pub fn into_vec(mut self) -> Vec<u8> {
55        self.flush_remaining();
56        self.buf
57    }
58
59    #[inline(always)]
60    fn ensure_capacity(&mut self) {
61        if self.pos + 8 > self.buf.capacity() {
62            primitives::set_vec_len(&mut self.buf, self.pos);
63            self.buf.reserve(64);
64        }
65    }
66
67    #[inline(always)]
68    pub fn write_bits(&mut self, value: u32, n: u8) {
69        if n == 0 {
70            return;
71        }
72        assert!(n <= 25, "bit count must be <= 25");
73        assert!(value < (1u32 << n), "value does not fit in bit count");
74        self.bits |= (value as u64) << self.bits_used;
75        self.bits_used += n;
76        if self.bits_used >= 32 {
77            self.ensure_capacity();
78            primitives::write_u64_le_unaligned(&mut self.buf, self.pos, self.bits);
79            let nb = (self.bits_used >> 3) as usize;
80            self.pos += nb;
81            self.bits >>= nb << 3;
82            self.bits_used &= 7;
83        }
84    }
85
86    pub fn flush_remaining(&mut self) {
87        primitives::set_vec_len(&mut self.buf, self.pos);
88        while self.bits_used > 0 {
89            self.buf.push(self.bits as u8);
90            self.bits >>= 8;
91            self.bits_used = self.bits_used.saturating_sub(8);
92        }
93        self.pos = self.buf.len();
94    }
95
96    pub fn close_reverse_stream(&mut self) {
97        self.write_bits(1, 1);
98        self.flush_remaining();
99    }
100
101    pub fn bits_written(&self) -> usize {
102        self.pos * 8 + self.bits_used as usize
103    }
104
105    pub fn into_bytes(mut self) -> Vec<u8> {
106        self.flush_remaining();
107        self.buf
108    }
109
110    pub fn as_bytes(&mut self) -> &[u8] {
111        primitives::set_vec_len(&mut self.buf, self.pos);
112        &self.buf
113    }
114
115    pub fn write_byte(&mut self, b: u8) {
116        self.write_bits(b as u32, 8);
117    }
118
119    pub fn write_bytes(&mut self, bytes: &[u8]) {
120        for &b in bytes {
121            self.write_byte(b);
122        }
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    #[should_panic(expected = "bit count must be <= 25")]
132    fn write_bits_rejects_large_public_bit_count() {
133        // Regression: `write_bits` is a safe public API, so violating its
134        // maximum bit-count contract must stop before the flush path shifts by
135        // an invalid amount.
136        let mut writer = BitWriter::new();
137        writer.write_bits(0, 255);
138    }
139
140    #[test]
141    fn write_and_read_back() {
142        let mut w = BitWriter::new();
143        w.write_bits(0b101, 3);
144        w.write_bits(0b11, 2);
145        w.write_bits(0b000, 3);
146        let bytes = w.into_bytes();
147        assert_eq!(bytes, [0b0001_1101]);
148    }
149
150    #[test]
151    fn write_cross_byte() {
152        let mut w = BitWriter::new();
153        w.write_bits(0xFF, 8);
154        w.write_bits(0x01, 8);
155        let bytes = w.into_bytes();
156        assert_eq!(bytes, [0xFF, 0x01]);
157    }
158
159    #[test]
160    fn write_partial() {
161        let mut w = BitWriter::new();
162        w.write_bits(0b1, 1);
163        let bytes = w.into_bytes();
164        assert_eq!(bytes, [0b0000_0001]);
165    }
166
167    #[test]
168    fn reverse_stream_roundtrip() {
169        use crate::bitstream::reader_reverse::ReverseBitReader;
170
171        let mut w = BitWriter::new();
172        w.write_bits(0b1010, 4);
173        w.write_bits(0b0101, 4);
174        w.close_reverse_stream();
175        let bytes = w.into_bytes();
176
177        let mut r = ReverseBitReader::new(&bytes).unwrap();
178        assert_eq!(r.read_bits(4).unwrap(), 0b0101);
179        assert_eq!(r.read_bits(4).unwrap(), 0b1010);
180        assert_eq!(r.bits_remaining(), 0);
181    }
182}
zrip_core/bitstream/writer.rs

zrip_core/bitstream/
writer.rs
