zrip_core/bitstream/
writer.rs

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

zrip_core/bitstream/
writer.rs
