|
1 | | -use std::io::{self, Read, Write}; |
2 | | - |
3 | | -use byteorder::{ReadBytesExt, WriteBytesExt}; |
4 | | - |
5 | | -/// A simple way to write individual bits to an input implementing [Write]. |
6 | | -pub struct BitWriter<'a, O: Write + WriteBytesExt> { |
7 | | - output: &'a mut O, |
8 | | - |
9 | | - current_byte: u8, |
10 | | - |
| 1 | +pub struct BitIo { |
| 2 | + data: Vec<u8>, |
11 | 3 | byte_offset: usize, |
12 | 4 | bit_offset: usize, |
13 | 5 |
|
14 | 6 | byte_size: usize, |
15 | 7 | } |
16 | 8 |
|
17 | | -impl<'a, O: Write + WriteBytesExt> BitWriter<'a, O> { |
18 | | - /// Create a new BitWriter wrapper around something which |
19 | | - /// implements [Write]. |
20 | | - pub fn new(output: &'a mut O) -> Self { |
| 9 | +impl BitIo { |
| 10 | + /// Create a new BitIO reader and writer over some data |
| 11 | + pub fn new(data: Vec<u8>) -> Self { |
21 | 12 | Self { |
22 | | - output, |
23 | | - |
24 | | - current_byte: 0, |
25 | | - |
| 13 | + data, |
26 | 14 | byte_offset: 0, |
27 | 15 | bit_offset: 0, |
28 | | - |
29 | 16 | byte_size: 0, |
30 | 17 | } |
31 | 18 | } |
32 | 19 |
|
33 | | - /// Get the number of whole bytes written to the stream. |
34 | | - pub fn byte_size(&self) -> usize { |
35 | | - self.byte_size |
36 | | - } |
37 | | - |
38 | | - /// Get the bit offset within the current byte. |
39 | | - pub fn bit_offset(&self) -> u8 { |
40 | | - self.bit_offset as u8 |
| 20 | + /// Get the byte offset of the reader |
| 21 | + pub fn byte_offset(&self) -> usize { |
| 22 | + self.byte_offset |
41 | 23 | } |
42 | 24 |
|
43 | | - /// Check if the stream is aligned to a byte. |
44 | | - pub fn aligned(&self) -> bool { |
45 | | - self.bit_offset() == 0 |
| 25 | + /// Get the byte size of the reader |
| 26 | + pub fn byte_size(&self) -> usize { |
| 27 | + self.byte_size |
46 | 28 | } |
47 | 29 |
|
48 | | - /// Align the writer to the nearest byte by padding with zero bits. |
49 | | - /// |
50 | | - /// Returns the number of zero bits |
51 | | - pub fn flush(&mut self) -> Result<usize, io::Error> { |
52 | | - self.byte_offset += 1; |
53 | | - |
54 | | - // Write out the current byte unfinished |
55 | | - self.output.write_u8(self.current_byte).unwrap(); |
56 | | - self.current_byte = 0; |
57 | | - self.bit_offset = 0; |
58 | | - |
59 | | - Ok(8 - self.bit_offset) |
| 30 | + /// Get the current bytes up to `byte_size` in the reader |
| 31 | + pub fn bytes(&self) -> Vec<u8> { |
| 32 | + self.data[..self.byte_size].to_vec() |
60 | 33 | } |
61 | 34 |
|
62 | | - /// Write some bits to the output. |
63 | | - pub fn write_bit(&mut self, data: u64, bit_len: usize) { |
64 | | - if bit_len > 64 { |
65 | | - panic!("Cannot write more than 64 bits at once."); |
66 | | - } else if bit_len == 0 { |
67 | | - panic!("Must write 1 or more bits.") |
| 35 | + /// Read some bits from the buffer |
| 36 | + pub fn read_bit(&mut self, bit_len: usize) -> u64 { |
| 37 | + if bit_len > 8 * 8 { |
| 38 | + panic!("Cannot read more than 64 bits") |
68 | 39 | } |
69 | 40 |
|
70 | 41 | if bit_len % 8 == 0 && self.bit_offset == 0 { |
71 | | - self.write(data, bit_len / 8); |
72 | | - return; |
| 42 | + return self.read(bit_len / 8); |
73 | 43 | } |
74 | 44 |
|
| 45 | + let mut result = 0; |
75 | 46 | for i in 0..bit_len { |
76 | | - let bit_value = (data >> i) & 1; |
77 | | - |
78 | | - self.current_byte &= !(1 << self.bit_offset); |
79 | | - |
80 | | - self.current_byte |= (bit_value << self.bit_offset) as u8; |
81 | | - |
| 47 | + let bit_value = ((self.data[self.byte_offset] as usize >> self.bit_offset) & 1) as u64; |
82 | 48 | self.bit_offset += 1; |
83 | | - if self.bit_offset >= 8 { |
| 49 | + |
| 50 | + if self.bit_offset == 8 { |
84 | 51 | self.byte_offset += 1; |
85 | 52 | self.bit_offset = 0; |
86 | | - |
87 | | - self.output.write_u8(self.current_byte).unwrap(); |
88 | | - self.current_byte = 0; |
89 | 53 | } |
| 54 | + |
| 55 | + result |= bit_value << i; |
90 | 56 | } |
91 | 57 |
|
92 | | - self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; |
| 58 | + result |
93 | 59 | } |
94 | 60 |
|
95 | | - /// Write some bytes to the output. |
96 | | - pub fn write(&mut self, data: u64, byte_len: usize) { |
| 61 | + /// Read some bytes from the buffer |
| 62 | + pub fn read(&mut self, byte_len: usize) -> u64 { |
97 | 63 | if byte_len > 8 { |
98 | | - panic!("Cannot write more than 8 bytes at once.") |
99 | | - } else if byte_len == 0 { |
100 | | - panic!("Must write 1 or more bytes.") |
| 64 | + panic!("Cannot read more than 8 bytes") |
101 | 65 | } |
102 | 66 |
|
103 | | - self.output |
104 | | - .write_all(&data.to_le_bytes()[..byte_len]) |
105 | | - .unwrap(); |
| 67 | + let mut padded_slice = [0u8; 8]; |
| 68 | + padded_slice.copy_from_slice(&self.data[self.byte_offset..self.byte_offset + byte_len]); |
106 | 69 | self.byte_offset += byte_len; |
107 | 70 |
|
108 | | - self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; |
| 71 | + u64::from_le_bytes(padded_slice) |
109 | 72 | } |
110 | | -} |
111 | | - |
112 | | -/// A simple way to read individual bits from an input implementing [Read]. |
113 | | -pub struct BitReader<'a, I: Read + ReadBytesExt> { |
114 | | - input: &'a mut I, |
115 | 73 |
|
116 | | - current_byte: Option<u8>, |
117 | | - |
118 | | - byte_offset: usize, |
119 | | - bit_offset: usize, |
120 | | -} |
121 | | - |
122 | | -impl<'a, I: Read + ReadBytesExt> BitReader<'a, I> { |
123 | | - /// Create a new BitReader wrapper around something which |
124 | | - /// implements [Write]. |
125 | | - pub fn new(input: &'a mut I) -> Self { |
126 | | - let first = input.read_u8().unwrap(); |
127 | | - Self { |
128 | | - input, |
129 | | - |
130 | | - current_byte: Some(first), |
131 | | - |
132 | | - byte_offset: 0, |
133 | | - bit_offset: 0, |
134 | | - } |
135 | | - } |
136 | | - |
137 | | - /// Get the number of whole bytes read from the stream. |
138 | | - pub fn byte_offset(&self) -> usize { |
139 | | - self.byte_offset |
140 | | - } |
141 | | - |
142 | | - /// Read some bits from the input. |
143 | | - pub fn read_bit(&mut self, bit_len: usize) -> u64 { |
144 | | - if bit_len > 64 { |
145 | | - panic!("Cannot read more than 64 bits at once.") |
146 | | - } else if bit_len == 0 { |
147 | | - panic!("Must read 1 or more bits.") |
| 74 | + /// Write some bits to the buffer |
| 75 | + pub fn write_bit(&mut self, data: u64, bit_len: usize) { |
| 76 | + if bit_len > 8 * 8 { |
| 77 | + panic!("Cannot write more than 64 bits"); |
148 | 78 | } |
149 | 79 |
|
150 | 80 | if bit_len % 8 == 0 && self.bit_offset == 0 { |
151 | | - return self.read(bit_len / 8); |
| 81 | + self.write(data, bit_len / 8); |
| 82 | + return; |
152 | 83 | } |
153 | 84 |
|
154 | | - let mut result = 0; |
155 | 85 | for i in 0..bit_len { |
156 | | - let bit_value = ((self.current_byte.unwrap() as usize >> self.bit_offset) & 1) as u64; |
157 | | - self.bit_offset += 1; |
| 86 | + let bit_value = (data >> i) & 1; |
| 87 | + |
| 88 | + self.data[self.byte_offset] &= !(1 << self.bit_offset); |
| 89 | + |
| 90 | + self.data[self.byte_offset] |= (bit_value << self.bit_offset) as u8; |
158 | 91 |
|
| 92 | + self.bit_offset += 1; |
159 | 93 | if self.bit_offset == 8 { |
160 | 94 | self.byte_offset += 1; |
161 | 95 | self.bit_offset = 0; |
162 | | - |
163 | | - self.current_byte = Some(self.input.read_u8().unwrap()); |
164 | 96 | } |
165 | | - |
166 | | - result |= bit_value << i; |
167 | 97 | } |
168 | 98 |
|
169 | | - result |
| 99 | + self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; |
170 | 100 | } |
171 | 101 |
|
172 | | - /// Read some bytes from the input. |
173 | | - pub fn read(&mut self, byte_len: usize) -> u64 { |
| 102 | + pub fn write(&mut self, data: u64, byte_len: usize) { |
174 | 103 | if byte_len > 8 { |
175 | | - panic!("Cannot read more than 8 bytes at once.") |
176 | | - } else if byte_len == 0 { |
177 | | - panic!("Must read 1 or more bytes") |
| 104 | + panic!("Cannot write more than 8 bytes") |
178 | 105 | } |
179 | 106 |
|
180 | | - let mut padded_slice = vec![0u8; byte_len]; |
181 | | - self.input.read_exact(&mut padded_slice).unwrap(); |
182 | | - self.byte_offset += byte_len; |
| 107 | + let mut padded_slice = [0u8; 8]; |
| 108 | + padded_slice.copy_from_slice(&data.to_le_bytes()); |
183 | 109 |
|
184 | | - let extra_length = padded_slice.len() - byte_len; |
185 | | - padded_slice.extend_from_slice(&vec![0u8; extra_length]); |
| 110 | + self.data[self.byte_offset..self.byte_offset + byte_len] |
| 111 | + .copy_from_slice(&padded_slice[..byte_len]); |
| 112 | + self.byte_offset += byte_len; |
186 | 113 |
|
187 | | - u64::from_le_bytes(padded_slice.try_into().unwrap()) |
| 114 | + self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; |
188 | 115 | } |
189 | 116 | } |
0 commit comments