use alloc::vec::Vec;
use crate::error::Error;
#[derive(Debug, Default)]
pub struct BitBuf {
pub buf: Vec<u8>,
pub byte_pos: usize,
pub bit_pos: u8,
pub last_byte: usize,
pub last_bit: u8,
}
impl BitBuf {
pub fn new() -> Self {
Self::default()
}
pub fn reset(&mut self, bytes: &[u8], last_bit: u8) {
self.buf.clear();
self.buf.extend_from_slice(bytes);
self.byte_pos = 0;
self.bit_pos = 0;
self.last_byte = self.buf.len().saturating_sub(1);
self.last_bit = last_bit;
}
pub fn at_end(&self) -> bool {
if self.buf.is_empty() {
return true;
}
if self.byte_pos > self.last_byte {
return true;
}
self.byte_pos == self.last_byte && self.bit_pos >= self.last_bit
}
pub fn bits_remaining(&self) -> u32 {
if self.buf.is_empty() {
return 0;
}
let end_bit = (self.last_byte as u64) * 8 + self.last_bit as u64;
let cur_bit = (self.byte_pos as u64) * 8 + self.bit_pos as u64;
end_bit.saturating_sub(cur_bit) as u32
}
pub fn peek16(&self) -> u16 {
let b0 = *self.buf.get(self.byte_pos).unwrap_or(&0) as u32;
let b1 = *self.buf.get(self.byte_pos + 1).unwrap_or(&0) as u32;
let b2 = *self.buf.get(self.byte_pos + 2).unwrap_or(&0) as u32;
let w = (b0 << 16) | (b1 << 8) | b2;
((w >> (8 - self.bit_pos as u32)) & 0xFFFF) as u16
}
pub fn peek_bits(&self, n: u32) -> u32 {
debug_assert!(n <= 32);
if n == 0 {
return 0;
}
let mut acc: u64 = 0;
for i in 0..5 {
let b = *self.buf.get(self.byte_pos + i).unwrap_or(&0) as u64;
acc = (acc << 8) | b;
}
let top = 40 - self.bit_pos as u32;
((acc >> (top - n)) & ((1u64 << n) - 1)) as u32
}
pub fn skip(&mut self, n: u32) {
let total = self.bit_pos as u32 + n;
self.byte_pos += (total >> 3) as usize;
self.bit_pos = (total & 7) as u8;
}
pub fn read(&mut self, n: u32) -> Result<u32, Error> {
debug_assert!(n <= 16);
if self.bits_remaining() < n {
return Err(Error::UnexpectedEnd);
}
let v = self.peek_bits(n);
self.skip(n);
Ok(v)
}
#[allow(dead_code)]
pub fn read1(&mut self) -> Result<u32, Error> {
self.read(1)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn msb_first_packed_reads() {
let mut br = BitBuf::new();
br.reset(&[0xB4, 0xE1], 8);
assert_eq!(br.read(4).unwrap(), 0b1011);
assert_eq!(br.read(4).unwrap(), 0b0100);
assert_eq!(br.read(8).unwrap(), 0xE1);
assert!(br.at_end());
}
#[test]
fn cross_byte_reads() {
let mut br = BitBuf::new();
br.reset(&[0x12, 0x34], 8);
assert_eq!(br.read(3).unwrap(), 0b000);
assert_eq!(br.read(9).unwrap(), 0b1_0010_0011);
assert_eq!(br.read(4).unwrap(), 0b0100);
}
#[test]
fn read_past_end_errors() {
let mut br = BitBuf::new();
br.reset(&[0xFF], 4);
assert_eq!(br.read(3).unwrap(), 0b111);
assert_eq!(br.read(2), Err(Error::UnexpectedEnd));
}
#[test]
fn peek_bits_independent_of_skip() {
let mut br = BitBuf::new();
br.reset(&[0xAB, 0xCD, 0xEF, 0x12], 8);
assert_eq!(br.peek_bits(16), 0xABCD);
br.skip(4);
assert_eq!(br.peek_bits(16), 0xBCDE);
}
}