#![warn(
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_possible_wrap
)]
pub struct BitReader<'a> {
data: &'a [u8],
pos: usize,
}
impl<'a> BitReader<'a> {
#[must_use]
pub fn new(data: &'a [u8]) -> Self {
BitReader { data, pos: 0 }
}
#[must_use]
pub fn bit_pos(&self) -> usize {
self.pos
}
pub fn seek_bits(&mut self, bit: usize) {
self.pos = bit;
}
pub fn read(&mut self, n: u32) -> u32 {
let v = self.peek(n);
self.pos += n as usize;
v
}
pub fn read_u8(&mut self, n: u32) -> u8 {
debug_assert!(n <= 8, "read_u8 of a {n}-bit field");
(self.read(n) & 0xFF) as u8
}
pub fn read_u16(&mut self, n: u32) -> u16 {
debug_assert!(n <= 16, "read_u16 of a {n}-bit field");
(self.read(n) & 0xFFFF) as u16
}
#[must_use]
pub fn peek(&self, n: u32) -> u32 {
if n == 0 {
return 0; }
let byte_idx = self.pos >> 3;
#[allow(clippy::cast_possible_truncation)] let bit_off = (self.pos & 7) as u32;
let word = if let Some(w) = self.data.get(byte_idx..byte_idx + 8) {
u64::from_be_bytes(w.try_into().unwrap())
} else {
let start = byte_idx.min(self.data.len());
let take = (self.data.len() - start).min(8);
let mut buf = [0u8; 8];
buf[..take].copy_from_slice(&self.data[start..start + take]);
u64::from_be_bytes(buf)
};
#[allow(clippy::cast_possible_truncation)]
let bits = ((word << bit_off) >> (64 - n)) as u32;
bits
}
pub fn skip(&mut self, n: u32) {
self.pos += n as usize;
}
pub fn read_bool(&mut self) -> bool {
self.read(1) != 0
}
}
#[derive(Default)]
pub struct BitWriter {
bytes: Vec<u8>,
nbits: u8,
cur: u8,
}
impl BitWriter {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_capacity(bytes: usize) -> Self {
Self {
bytes: Vec::with_capacity(bytes),
nbits: 0,
cur: 0,
}
}
pub fn write(&mut self, v: u32, n: u32) {
for i in (0..n).rev() {
let bit = ((v >> i) & 1) as u8;
self.cur = (self.cur << 1) | bit;
self.nbits += 1;
if self.nbits == 8 {
self.bytes.push(self.cur);
self.cur = 0;
self.nbits = 0;
}
}
}
#[must_use]
pub fn bit_len(&self) -> usize {
self.bytes.len() * 8 + self.nbits as usize
}
#[must_use]
pub fn finish(mut self) -> Vec<u8> {
if self.nbits > 0 {
self.cur <<= 8 - self.nbits;
self.bytes.push(self.cur);
}
self.bytes
}
}
#[cfg(test)]
mod peek_tests {
use super::*;
fn peek_bitwise(data: &[u8], pos: usize, n: u32) -> u32 {
let mut v = 0u32;
for p in pos..pos + n as usize {
let byte = data.get(p >> 3).copied().unwrap_or(0);
let bit = (byte >> (7 - (p & 7))) & 1;
v = (v << 1) | u32::from(bit);
}
v
}
#[test]
fn peek_matches_bitwise_reference() {
let mut s = 0x2545_F491u32;
let data: Vec<u8> = (0..64)
.map(|_| {
s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
(s >> 24) as u8
})
.collect();
for pos in 0..(data.len() * 8 + 40) {
for n in 0..=32u32 {
let r = BitReader { data: &data, pos };
assert_eq!(
r.peek(n),
peek_bitwise(&data, pos, n),
"mismatch at pos={pos} n={n}"
);
}
}
}
#[test]
fn peek_on_empty_buffer_is_zero() {
let r = BitReader { data: &[], pos: 0 };
for n in 0..=32u32 {
assert_eq!(r.peek(n), 0);
}
}
}