use std::io;
const STORE_SIZE: u8 = 128;
#[derive(Debug)]
pub struct BitWriter<W> {
pub data: W,
pub bit_store: u128,
pub bit_count: u8,
}
impl<W: io::Write> BitWriter<W> {
pub const fn new(data: W) -> Self {
Self {
data,
bit_store: 0,
bit_count: 0,
}
}
pub fn write_bits(&mut self, bits: u16, count: u8) -> io::Result<()> {
if self.bit_count + count >= STORE_SIZE {
self.flush_bytes()?;
}
let bits = u128::from(bits);
let mask = (1 << count) - 1;
self.bit_store |= (bits & mask) << self.bit_count;
self.bit_count += count;
Ok(())
}
pub fn force_align(&mut self) -> io::Result<()> {
let remainder = self.bit_count % 8;
if remainder != 0 {
self.bit_count += 8 - remainder;
}
self.flush_bytes()
}
fn flush_bytes(&mut self) -> io::Result<()> {
let full_byte_count = self.bit_count.saturating_div(8);
let buf = self.bit_store.to_le_bytes();
self.data
.write_all(&buf.as_slice()[..full_byte_count as usize])?;
self.bit_store >>= full_byte_count.saturating_mul(8);
self.bit_count -= full_byte_count.saturating_mul(8);
Ok(())
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
#[test]
fn test_write_basic_bits() {
let mut bw = BitWriter::new(Vec::new());
bw.write_bits(0b1010, 4).unwrap();
assert_eq!(bw.bit_store, 0b1010);
assert_eq!(bw.bit_count, 4);
bw.write_bits(0b1100, 4).unwrap();
assert_eq!(bw.bit_store, 0b1100_1010);
assert_eq!(bw.bit_count, 8);
}
#[test]
fn test_flush_exact_byte() {
let mut bw = BitWriter::new(Vec::new());
bw.write_bits(0b1100_1010, 8).unwrap();
bw.flush_bytes().unwrap();
assert_eq!(bw.data, vec![0b1100_1010]);
assert_eq!(bw.bit_count, 0);
assert_eq!(bw.bit_store, 0);
}
#[test]
fn test_cross_byte_boundary() {
let mut bw = BitWriter::new(Vec::new());
bw.write_bits(0x533, 12).unwrap();
bw.flush_bytes().unwrap();
assert_eq!(bw.data, vec![0x33]);
assert_eq!(bw.bit_count, 4);
assert_eq!(bw.bit_store, 0x05);
}
#[test]
fn test_force_align() {
let mut bw = BitWriter::new(Vec::new());
bw.write_bits(0b101, 3).unwrap();
bw.force_align().unwrap();
assert_eq!(bw.data, vec![0x05]);
assert_eq!(bw.bit_count, 0);
assert_eq!(bw.bit_store, 0);
}
}