1use crate::{
2 constants::{MTU_SIZE_BITS, MTU_SIZE_BYTES},
3 BitCounter, OutgoingPacket, OwnedBitReader,
4};
5
6pub 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
15pub struct BitWriter {
17 scratch: u8,
18 scratch_index: u8,
19 buffer: [u8; MTU_SIZE_BYTES],
20 buffer_index: usize,
21 current_bits: u32,
22 max_bits: u32,
23}
24
25impl BitWriter {
26 #[allow(clippy::new_without_default)]
27 pub fn new() -> Self {
28 Self {
29 scratch: 0,
30 scratch_index: 0,
31 buffer: [0; MTU_SIZE_BYTES],
32 buffer_index: 0,
33 current_bits: 0,
34 max_bits: MTU_SIZE_BITS,
35 }
36 }
37
38 pub fn with_capacity(bit_capacity: u32) -> Self {
39 Self {
40 scratch: 0,
41 scratch_index: 0,
42 buffer: [0; MTU_SIZE_BYTES],
43 buffer_index: 0,
44 current_bits: 0,
45 max_bits: bit_capacity,
46 }
47 }
48
49 pub fn with_max_capacity() -> Self {
50 Self::with_capacity(u32::MAX)
51 }
52
53 fn finalize(&mut self) {
54 if self.scratch_index > 0 {
55 self.buffer[self.buffer_index] =
56 (self.scratch << (8 - self.scratch_index)).reverse_bits();
57 self.buffer_index += 1;
58 }
59 self.max_bits = 0;
60 }
61
62 pub fn to_packet(mut self) -> OutgoingPacket {
63 self.finalize();
64 OutgoingPacket::new(self.buffer_index, self.buffer)
65 }
66
67 pub fn to_owned_reader(mut self) -> OwnedBitReader {
68 self.finalize();
69 OwnedBitReader::new(&self.buffer[0..self.buffer_index])
70 }
71
72 pub fn to_bytes(mut self) -> Box<[u8]> {
73 self.finalize();
74 Box::from(&self.buffer[0..self.buffer_index])
75 }
76
77 pub fn counter(&self) -> BitCounter {
78 return BitCounter::new(self.current_bits, self.current_bits, self.max_bits);
79 }
80
81 pub fn reserve_bits(&mut self, bits: u32) {
82 self.max_bits -= bits;
83 }
84
85 pub fn release_bits(&mut self, bits: u32) {
86 self.max_bits += bits;
87 }
88
89 pub fn bits_free(&self) -> u32 {
90 self.max_bits - self.current_bits
91 }
92}
93
94impl BitWrite for BitWriter {
95 fn write_bit(&mut self, bit: bool) {
96 if self.current_bits >= self.max_bits {
97 panic!("Write overflow!");
98 }
99 self.scratch <<= 1;
100
101 if bit {
102 self.scratch |= 1;
103 }
104
105 self.scratch_index += 1;
106 self.current_bits += 1;
107
108 if self.scratch_index >= 8 {
109 self.buffer[self.buffer_index] = self.scratch.reverse_bits();
110
111 self.buffer_index += 1;
112 self.scratch_index -= 8;
113 self.scratch = 0;
114 }
115 }
116
117 fn write_byte(&mut self, byte: u8) {
118 let mut temp = byte;
119 for _ in 0..8 {
120 self.write_bit(temp & 1 != 0);
121 temp >>= 1;
122 }
123 }
124
125 fn count_bits(&mut self, _: u32) {
126 panic!("This method should not be called for BitWriter!");
127 }
128
129 fn is_counter(&self) -> bool {
130 false
131 }
132}
133
134mod tests {
135
136 #[test]
137 fn read_write_1_bit() {
138 use crate::{
139 bit_reader::BitReader,
140 bit_writer::{BitWrite, BitWriter},
141 };
142
143 let mut writer = BitWriter::new();
144
145 writer.write_bit(true);
146
147 let buffer = writer.to_bytes();
148
149 let mut reader = BitReader::new(&buffer);
150
151 assert!(reader.read_bit().unwrap());
152 }
153
154 #[test]
155 fn read_write_3_bits() {
156 use crate::{
157 bit_reader::BitReader,
158 bit_writer::{BitWrite, BitWriter},
159 };
160
161 let mut writer = BitWriter::new();
162
163 writer.write_bit(false);
164 writer.write_bit(true);
165 writer.write_bit(true);
166
167 let buffer = writer.to_bytes();
168
169 let mut reader = BitReader::new(&buffer);
170
171 assert!(!reader.read_bit().unwrap());
172 assert!(reader.read_bit().unwrap());
173 assert!(reader.read_bit().unwrap());
174 }
175
176 #[test]
177 fn read_write_8_bits() {
178 use crate::{
179 bit_reader::BitReader,
180 bit_writer::{BitWrite, BitWriter},
181 };
182
183 let mut writer = BitWriter::new();
184
185 writer.write_bit(false);
186 writer.write_bit(true);
187 writer.write_bit(false);
188 writer.write_bit(true);
189
190 writer.write_bit(true);
191 writer.write_bit(false);
192 writer.write_bit(false);
193 writer.write_bit(false);
194
195 let buffer = writer.to_bytes();
196
197 let mut reader = BitReader::new(&buffer);
198
199 assert!(!reader.read_bit().unwrap());
200 assert!(reader.read_bit().unwrap());
201 assert!(!reader.read_bit().unwrap());
202 assert!(reader.read_bit().unwrap());
203
204 assert!(reader.read_bit().unwrap());
205 assert!(!reader.read_bit().unwrap());
206 assert!(!reader.read_bit().unwrap());
207 assert!(!reader.read_bit().unwrap());
208 }
209
210 #[test]
211 fn read_write_13_bits() {
212 use crate::{
213 bit_reader::BitReader,
214 bit_writer::{BitWrite, BitWriter},
215 };
216
217 let mut writer = BitWriter::new();
218
219 writer.write_bit(false);
220 writer.write_bit(true);
221 writer.write_bit(false);
222 writer.write_bit(true);
223
224 writer.write_bit(true);
225 writer.write_bit(false);
226 writer.write_bit(false);
227 writer.write_bit(false);
228
229 writer.write_bit(true);
230 writer.write_bit(false);
231 writer.write_bit(true);
232 writer.write_bit(true);
233
234 writer.write_bit(true);
235
236 let buffer = writer.to_bytes();
237
238 let mut reader = BitReader::new(&buffer);
239
240 assert!(!reader.read_bit().unwrap());
241 assert!(reader.read_bit().unwrap());
242 assert!(!reader.read_bit().unwrap());
243 assert!(reader.read_bit().unwrap());
244
245 assert!(reader.read_bit().unwrap());
246 assert!(!reader.read_bit().unwrap());
247 assert!(!reader.read_bit().unwrap());
248 assert!(!reader.read_bit().unwrap());
249
250 assert!(reader.read_bit().unwrap());
251 assert!(!reader.read_bit().unwrap());
252 assert!(reader.read_bit().unwrap());
253 assert!(reader.read_bit().unwrap());
254
255 assert!(reader.read_bit().unwrap());
256 }
257
258 #[test]
259 fn read_write_16_bits() {
260 use crate::{
261 bit_reader::BitReader,
262 bit_writer::{BitWrite, BitWriter},
263 };
264
265 let mut writer = BitWriter::new();
266
267 writer.write_bit(false);
268 writer.write_bit(true);
269 writer.write_bit(false);
270 writer.write_bit(true);
271
272 writer.write_bit(true);
273 writer.write_bit(false);
274 writer.write_bit(false);
275 writer.write_bit(false);
276
277 writer.write_bit(true);
278 writer.write_bit(false);
279 writer.write_bit(true);
280 writer.write_bit(true);
281
282 writer.write_bit(true);
283 writer.write_bit(false);
284 writer.write_bit(true);
285 writer.write_bit(true);
286
287 let buffer = writer.to_bytes();
288
289 let mut reader = BitReader::new(&buffer);
290
291 assert!(!reader.read_bit().unwrap());
292 assert!(reader.read_bit().unwrap());
293 assert!(!reader.read_bit().unwrap());
294 assert!(reader.read_bit().unwrap());
295
296 assert!(reader.read_bit().unwrap());
297 assert!(!reader.read_bit().unwrap());
298 assert!(!reader.read_bit().unwrap());
299 assert!(!reader.read_bit().unwrap());
300
301 assert!(reader.read_bit().unwrap());
302 assert!(!reader.read_bit().unwrap());
303 assert!(reader.read_bit().unwrap());
304 assert!(reader.read_bit().unwrap());
305
306 assert!(reader.read_bit().unwrap());
307 assert!(!reader.read_bit().unwrap());
308 assert!(reader.read_bit().unwrap());
309 assert!(reader.read_bit().unwrap());
310 }
311
312 #[test]
313 fn read_write_1_byte() {
314 use crate::{
315 bit_reader::BitReader,
316 bit_writer::{BitWrite, BitWriter},
317 };
318
319 let mut writer = BitWriter::new();
320
321 writer.write_byte(123);
322
323 let buffer = writer.to_bytes();
324
325 let mut reader = BitReader::new(&buffer);
326
327 assert_eq!(123, reader.read_byte().unwrap());
328 }
329
330 #[test]
331 fn read_write_5_bytes() {
332 use crate::{
333 bit_reader::BitReader,
334 bit_writer::{BitWrite, BitWriter},
335 };
336
337 let mut writer = BitWriter::new();
338
339 writer.write_byte(48);
340 writer.write_byte(151);
341 writer.write_byte(62);
342 writer.write_byte(34);
343 writer.write_byte(2);
344
345 let buffer = writer.to_bytes();
346
347 let mut reader = BitReader::new(&buffer);
348
349 assert_eq!(48, reader.read_byte().unwrap());
350 assert_eq!(151, reader.read_byte().unwrap());
351 assert_eq!(62, reader.read_byte().unwrap());
352 assert_eq!(34, reader.read_byte().unwrap());
353 assert_eq!(2, reader.read_byte().unwrap());
354 }
355}