1use crate::error::{Error, Result};
4
5pub fn read_bit(data: &[u8], bit_offset: usize) -> Result<bool> {
6 let byte_index = bit_offset / 8;
7 let bit_index = bit_offset % 8;
8 let byte = *data.get(byte_index).ok_or(Error::Truncated {
9 offset: byte_index as u64,
10 })?;
11 Ok(((byte >> (7 - bit_index)) & 1) != 0)
12}
13
14#[derive(Debug, Clone, Copy)]
16pub struct BitReader<'a> {
17 data: &'a [u8],
18 bit_offset: usize,
19}
20
21impl<'a> BitReader<'a> {
22 pub const fn new(data: &'a [u8]) -> Self {
23 Self {
24 data,
25 bit_offset: 0,
26 }
27 }
28
29 pub const fn with_offset(data: &'a [u8], bit_offset: usize) -> Self {
30 Self { data, bit_offset }
31 }
32
33 pub const fn bit_offset(&self) -> usize {
34 self.bit_offset
35 }
36
37 pub fn read(&mut self, bit_count: usize) -> Result<u64> {
38 if bit_count == 0 {
39 return Ok(0);
40 }
41 require_u64_width(bit_count)?;
42 let end_bit_offset = self
43 .bit_offset
44 .checked_add(bit_count)
45 .ok_or_else(|| Error::Other("bit offset overflow".into()))?;
46
47 let mut remaining = bit_count;
48 let mut value = 0u64;
49
50 while remaining > 0 {
51 let byte_index = self.bit_offset / 8;
52 let bit_index = self.bit_offset % 8;
53 let byte = *self.data.get(byte_index).ok_or(Error::Truncated {
54 offset: byte_index as u64,
55 })?;
56 let available = 8 - bit_index;
57 let take = remaining.min(available);
58 let mask = ((1u16 << take) - 1) as u8;
59 let shift = available - take;
60 let bits = (byte >> shift) & mask;
61
62 value = (value << take) | u64::from(bits);
63 self.bit_offset = self
64 .bit_offset
65 .checked_add(take)
66 .ok_or_else(|| Error::Other("bit offset overflow".into()))?;
67 remaining -= take;
68 }
69
70 debug_assert_eq!(self.bit_offset, end_bit_offset);
71 Ok(value)
72 }
73
74 pub fn read_bool(&mut self) -> Result<bool> {
75 Ok(self.read(1)? != 0)
76 }
77
78 pub fn read_signed(&mut self, bit_count: usize) -> Result<i64> {
79 if bit_count == 0 {
80 return Ok(0);
81 }
82 require_u64_width(bit_count)?;
83
84 let value = self.read(bit_count)?;
85 let sign_mask = 1u64 << (bit_count - 1);
86 if value & sign_mask == 0 {
87 return i64::try_from(value)
88 .map_err(|_| Error::Other("signed value exceeds i64 range".into()));
89 }
90
91 let magnitude_mask = sign_mask - 1;
92 let magnitude = value & magnitude_mask;
93 let magnitude = i64::try_from(magnitude)
94 .map_err(|_| Error::Other("signed value exceeds i64 range".into()))?;
95 Ok(-magnitude)
96 }
97}
98
99#[derive(Debug, Clone, Default)]
101pub struct BitWriter {
102 bytes: Vec<u8>,
103 bit_offset: usize,
104}
105
106impl BitWriter {
107 pub const fn new() -> Self {
108 Self {
109 bytes: Vec::new(),
110 bit_offset: 0,
111 }
112 }
113
114 pub fn with_capacity_bits(bit_capacity: usize) -> Self {
115 Self {
116 bytes: Vec::with_capacity(bit_capacity.div_ceil(8)),
117 bit_offset: 0,
118 }
119 }
120
121 pub fn bit_len(&self) -> usize {
122 self.bit_offset
123 }
124
125 pub fn byte_len(&self) -> usize {
126 self.bytes.len()
127 }
128
129 pub fn is_empty(&self) -> bool {
130 self.bit_offset == 0
131 }
132
133 pub fn as_bytes(&self) -> &[u8] {
134 &self.bytes
135 }
136
137 pub fn into_bytes(self) -> Vec<u8> {
138 self.bytes
139 }
140
141 pub fn write(&mut self, value: u64, bit_count: usize) -> Result<()> {
142 if bit_count == 0 {
143 if value == 0 {
144 return Ok(());
145 }
146 return Err(Error::Other(
147 "non-zero value cannot be written with zero bits".into(),
148 ));
149 }
150 require_u64_width(bit_count)?;
151 if bit_count < u64::BITS as usize && (value >> bit_count) != 0 {
152 return Err(Error::Other(format!(
153 "value {value} does not fit in {bit_count} bits"
154 )));
155 }
156
157 let start_bit_offset = self.bit_offset;
158 let end_bit_offset = self
159 .bit_offset
160 .checked_add(bit_count)
161 .ok_or_else(|| Error::Other("bit offset overflow".into()))?;
162 let required_bytes = end_bit_offset.div_ceil(8);
163 if required_bytes > self.bytes.len() {
164 let additional = required_bytes - self.bytes.len();
165 self.bytes
166 .try_reserve(additional)
167 .map_err(|e| Error::Other(format!("failed to reserve {additional} bytes: {e}")))?;
168 self.bytes.resize(required_bytes, 0);
169 }
170
171 for target_index in 0..bit_count {
172 let source_shift = bit_count - 1 - target_index;
173 let bit = ((value >> source_shift) & 1) as u8;
174 if bit != 0 {
175 let target_offset = start_bit_offset + target_index;
176 let byte_index = target_offset / 8;
177 let shift = 7 - (target_offset % 8);
178 self.bytes[byte_index] |= 1 << shift;
179 }
180 }
181
182 self.bit_offset = end_bit_offset;
183 Ok(())
184 }
185
186 pub fn align_to_byte(&mut self) -> Result<()> {
187 let remainder = self.bit_offset % 8;
188 if remainder != 0 {
189 self.bit_offset = self
190 .bit_offset
191 .checked_add(8 - remainder)
192 .ok_or_else(|| Error::Other("bit offset overflow".into()))?;
193 }
194 Ok(())
195 }
196}
197
198fn require_u64_width(bit_count: usize) -> Result<()> {
199 if bit_count <= u64::BITS as usize {
200 return Ok(());
201 }
202
203 Err(Error::UnsupportedPackingWidth(
204 u8::try_from(bit_count).unwrap_or(u8::MAX),
205 ))
206}
207
208#[cfg(test)]
209mod tests {
210 use super::{read_bit, BitReader, BitWriter};
211
212 #[test]
213 fn reads_msb_first_across_byte_boundaries() {
214 let mut reader = BitReader::new(&[0b1011_0010, 0b0110_0000]);
215
216 assert_eq!(reader.read(3).unwrap(), 0b101);
217 assert_eq!(reader.read(5).unwrap(), 0b10010);
218 assert_eq!(reader.read(4).unwrap(), 0b0110);
219 assert_eq!(reader.bit_offset(), 12);
220 }
221
222 #[test]
223 fn reads_single_bits_by_offset() {
224 assert!(read_bit(&[0b1010_0000], 0).unwrap());
225 assert!(!read_bit(&[0b1010_0000], 1).unwrap());
226 assert!(read_bit(&[0b1010_0000], 2).unwrap());
227 assert!(read_bit(&[], 0).is_err());
228 }
229
230 #[test]
231 fn reads_grib_style_signed_magnitudes() {
232 let mut reader = BitReader::new(&[0b1000_0101]);
233 assert_eq!(reader.read_signed(8).unwrap(), -5);
234 }
235
236 #[test]
237 fn rejects_invalid_read_widths_without_panicking() {
238 let mut reader = BitReader::new(&[0xff; 9]);
239 assert!(reader.read(65).is_err());
240 assert!(reader.read_signed(65).is_err());
241 }
242
243 #[test]
244 fn reads_bits_written_by_writer() {
245 let mut writer = BitWriter::new();
246 writer.write(0b101, 3).unwrap();
247 writer.write(0b1111_0000, 8).unwrap();
248
249 let bytes = writer.into_bytes();
250 let mut reader = BitReader::new(&bytes);
251 assert_eq!(reader.read(3).unwrap(), 0b101);
252 assert_eq!(reader.read(8).unwrap(), 0b1111_0000);
253 }
254
255 #[test]
256 fn writes_msb_first_and_pads_final_byte() {
257 let mut writer = BitWriter::new();
258 writer.write(0b101, 3).unwrap();
259 writer.write(0b10010, 5).unwrap();
260 writer.write(0b0110, 4).unwrap();
261
262 assert_eq!(writer.bit_len(), 12);
263 assert_eq!(writer.as_bytes(), &[0b1011_0010, 0b0110_0000]);
264 }
265
266 #[test]
267 fn aligns_to_byte_and_tracks_lengths() {
268 let mut writer = BitWriter::new();
269 writer.write(0b101, 3).unwrap();
270 assert_eq!(writer.bit_len(), 3);
271 assert_eq!(writer.byte_len(), 1);
272
273 writer.align_to_byte().unwrap();
274 assert_eq!(writer.bit_len(), 8);
275 assert_eq!(writer.byte_len(), 1);
276
277 writer.write(0xff, 8).unwrap();
278 assert_eq!(writer.bit_len(), 16);
279 assert_eq!(writer.byte_len(), 2);
280 assert_eq!(writer.as_bytes(), &[0b1010_0000, 0xff]);
281 }
282
283 #[test]
284 fn rejects_values_that_do_not_fit_width() {
285 let mut writer = BitWriter::new();
286 assert!(writer.write(0b100, 2).is_err());
287 assert!(writer.write(0, 65).is_err());
288 }
289}