]> git.nihav.org Git - nihav.git/blob - nihav-core/src/io/bitwriter.rs
io/intcode: add variable-length code writing
[nihav.git] / nihav-core / src / io / bitwriter.rs
1 //! Bitstream writer functionality.
2 //!
3 //! Bitstream writer works on `Vec<u8>` and allows to write bits in different modes.
4 //!
5 //! # Examples
6 //!
7 //! Writing 17-bit value:
8 //! ```
9 //! use nihav_core::io::bitwriter::{BitWriter,BitWriterMode};
10 //!
11 //! # fn foo() -> Vec<u8> {
12 //! let mut bw = BitWriter::new(Vec::new(), BitWriterMode::BE);
13 //! bw.write(42, 17);
14 //! # bw.end()
15 //! # }
16 //! ```
17
18 /// Bitstream writing modes.
19 #[derive(Debug,Clone,Copy,PartialEq)]
20 pub enum BitWriterMode {
21 /// The stream is big endian MSB first.
22 BE,
23 /// The stream is little endian LSB first.
24 LE,
25 /// The stream is packed into 16-bit little-endian words MSB first.
26 LE16MSB,
27 /// The stream is packed into 32-bit little-endian words MSB first.
28 LE32MSB,
29 }
30
31 impl BitWriterMode {
32 fn is_be(self) -> bool { self != BitWriterMode::LE }
33 }
34
35 /// Bitstream writer.
36 pub struct BitWriter {
37 dst: Vec<u8>,
38 bitbuf: u32,
39 bits: u8,
40 start: usize,
41 mode: BitWriterMode,
42 }
43
44 impl BitWriter {
45 /// Creates a new instance of `BitWriter` that will append data to the input vector.
46 ///
47 /// # Examples
48 ///
49 /// ```
50 /// use nihav_core::io::bitwriter::{BitWriter,BitWriterMode};
51 ///
52 /// let mut bw = BitWriter::new(Vec::with_capacity(100), BitWriterMode::BE);
53 /// ```
54 pub fn new(dst: Vec<u8>, mode: BitWriterMode) -> Self {
55 let start = dst.len();
56 Self {
57 dst,
58 mode,
59 start,
60 bitbuf: 0,
61 bits: 0,
62 }
63 }
64 /// Writes single zero bit to the output.
65 pub fn write0(&mut self) { self.write_bit(false); }
66 /// Writes single set bit to the output.
67 pub fn write1(&mut self) { self.write_bit(true); }
68 /// Writes single bit.
69 pub fn write_bit(&mut self, bit: bool) {
70 if self.mode.is_be() {
71 self.bitbuf |= (bit as u32) << (31 - self.bits);
72 } else {
73 self.bitbuf |= (bit as u32) << self.bits;
74 }
75 self.bits += 1;
76 self.flush();
77 }
78 /// Writes `bits` bits of `val` value to the output.
79 #[allow(clippy::collapsible_if)]
80 pub fn write(&mut self, val: u32, bits: u8) {
81 if bits == 0 {
82 return;
83 }
84 if self.mode.is_be() {
85 if self.bits + bits <= 32 {
86 self.bitbuf |= val << (32 - self.bits - bits);
87 self.bits += bits;
88 self.flush();
89 } else {
90 let cbits = 32 - self.bits;
91 let bits2 = bits - cbits;
92 self.write(val >> bits2, cbits);
93 self.write(val & ((1 << bits2) - 1), bits2);
94 }
95 } else {
96 if self.bits + bits <= 32 {
97 self.bitbuf |= val << self.bits;
98 self.bits += bits;
99 self.flush();
100 } else {
101 let cbits = 32 - self.bits;
102 let bits2 = bits - cbits;
103 self.write(val & ((1 << cbits) - 1), cbits);
104 self.write(val >> cbits, bits2);
105 }
106 }
107 }
108 /// Writes `bits` bits of signed `val` value to the output.
109 pub fn write_s(&mut self, val: i32, bits: u8) {
110 self.write((val as u32) & ((1 << bits) - 1), bits);
111 }
112 /// Tells the amount of bits written so far.
113 pub fn tell(&self) -> usize {
114 (self.dst.len() - self.start) * 8 + (self.bits as usize)
115 }
116 fn flush(&mut self) {
117 match self.mode {
118 BitWriterMode::BE => {
119 while self.bits >= 8 {
120 self.dst.push((self.bitbuf >> 24) as u8);
121 self.bitbuf <<= 8;
122 self.bits -= 8;
123 }
124 },
125 BitWriterMode::LE => {
126 while self.bits >= 8 {
127 self.dst.push(self.bitbuf as u8);
128 self.bitbuf >>= 8;
129 self.bits -= 8;
130 }
131 },
132 BitWriterMode::LE16MSB => {
133 while self.bits >= 16 {
134 self.dst.push((self.bitbuf >> 16) as u8);
135 self.dst.push((self.bitbuf >> 24) as u8);
136 self.bitbuf <<= 16;
137 self.bits -= 16;
138 }
139 },
140 BitWriterMode::LE32MSB => {
141 if self.bits == 32 {
142 self.dst.push( self.bitbuf as u8);
143 self.dst.push((self.bitbuf >> 8) as u8);
144 self.dst.push((self.bitbuf >> 16) as u8);
145 self.dst.push((self.bitbuf >> 24) as u8);
146 self.bitbuf = 0;
147 self.bits = 0;
148 }
149 },
150 };
151 }
152 /// Finalises operations and returns the vector containing output data.
153 pub fn end(mut self) -> Vec<u8> {
154 self.flush();
155 if self.bits > 0 {
156 match self.mode {
157 BitWriterMode::BE => {
158 self.dst.push((self.bitbuf >> 24) as u8);
159 },
160 BitWriterMode::LE => {
161 self.dst.push(self.bitbuf as u8);
162 },
163 BitWriterMode::LE16MSB => {
164 self.dst.push((self.bitbuf >> 16) as u8);
165 self.dst.push((self.bitbuf >> 24) as u8);
166 },
167 BitWriterMode::LE32MSB => {
168 self.dst.push( self.bitbuf as u8);
169 self.dst.push((self.bitbuf >> 8) as u8);
170 self.dst.push((self.bitbuf >> 16) as u8);
171 self.dst.push((self.bitbuf >> 24) as u8);
172 },
173 };
174 }
175 self.dst
176 }
177 }
178
179 #[cfg(test)]
180 mod test {
181 use super::*;
182
183 #[test]
184 fn bw_works() {
185 let mut bw = BitWriter::new(Vec::new(), BitWriterMode::BE);
186 bw.write(43, 9);
187 let data = bw.end();
188 assert_eq!(&data, &[21, 128]);
189
190 let mut bw = BitWriter::new(Vec::new(), BitWriterMode::LE);
191 bw.write1();
192 bw.write(43, 9);
193 let data = bw.end();
194 assert_eq!(&data, &[87, 0]);
195
196 let mut bw = BitWriter::new(Vec::new(), BitWriterMode::LE32MSB);
197 bw.write(42, 9);
198 bw.write(42, 9);
199 let data = bw.end();
200 assert_eq!(&data, &[0, 128, 10, 21]);
201 }
202 }