#[derive(Debug, Clone)]
pub struct BitWriter {
data: Vec<u8>,
bit_pos: usize,
}
impl BitWriter {
pub fn new() -> Self {
Self {
data: Vec::new(),
bit_pos: 0,
}
}
pub fn write_bits(&mut self, value: u64, num_bits: usize) {
for i in 0..num_bits {
let bit = (value >> (num_bits - 1 - i)) & 1;
self.write_bit(bit != 0);
}
}
pub fn write_bit(&mut self, bit: bool) {
let byte_pos = self.bit_pos / 8;
let bit_offset = 7 - (self.bit_pos % 8);
while self.data.len() <= byte_pos {
self.data.push(0);
}
if bit {
self.data[byte_pos] |= 1 << bit_offset;
}
self.bit_pos += 1;
}
pub fn write_bits_to(data: &mut [u8], start_bit: usize, value: u64, num_bits: usize) {
for i in 0..num_bits {
let bit_pos = start_bit + i;
let byte_pos = bit_pos / 8;
let bit_offset = 7 - (bit_pos % 8);
if byte_pos >= data.len() {
break;
}
let bit = (value >> (num_bits - 1 - i)) & 1;
if bit != 0 {
data[byte_pos] |= 1 << bit_offset;
} else {
data[byte_pos] &= !(1 << bit_offset);
}
}
}
pub fn finish(mut self) -> Vec<u8> {
if self.bit_pos % 8 != 0 {
let byte_pos = self.bit_pos / 8;
while self.data.len() <= byte_pos {
self.data.push(0);
}
}
self.data
}
pub fn bit_position(&self) -> usize {
self.bit_pos
}
pub fn data(&self) -> &[u8] {
&self.data
}
}
impl Default for BitWriter {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct BitReader<'a> {
data: &'a [u8],
bit_pos: usize,
}
impl<'a> BitReader<'a> {
pub fn new(data: &'a [u8]) -> Self {
Self { data, bit_pos: 0 }
}
pub fn read_bits(&mut self, num_bits: usize) -> Result<u64, String> {
let mut value = 0u64;
for i in 0..num_bits {
let bit = self.read_bit()?;
if bit {
value |= 1 << (num_bits - 1 - i);
}
}
Ok(value)
}
pub fn read_bit(&mut self) -> Result<bool, String> {
let byte_pos = self.bit_pos / 8;
let bit_offset = 7 - (self.bit_pos % 8);
if byte_pos >= self.data.len() {
return Err("Unexpected end of data".to_string());
}
let bit = (self.data[byte_pos] >> bit_offset) & 1;
self.bit_pos += 1;
Ok(bit != 0)
}
pub fn read_bits_at(data: &[u8], start_bit: usize, num_bits: usize) -> u64 {
let mut value = 0u64;
for i in 0..num_bits {
let bit_pos = start_bit + i;
let byte_pos = bit_pos / 8;
let bit_offset = 7 - (bit_pos % 8);
if byte_pos >= data.len() {
break;
}
let bit = (data[byte_pos] >> bit_offset) & 1;
if bit != 0 {
value |= 1 << (num_bits - 1 - i);
}
}
value
}
pub fn bit_position(&self) -> usize {
self.bit_pos
}
pub fn has_bits(&self) -> bool {
self.bit_pos < self.data.len() * 8
}
pub fn reset(&mut self) {
self.bit_pos = 0;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bit_writer_basic() {
let mut writer = BitWriter::new();
writer.write_bits(0b1010, 4);
let data = writer.finish();
assert_eq!(data, vec![0b1010_0000]);
}
#[test]
fn test_bit_writer_multiple_bytes() {
let mut writer = BitWriter::new();
writer.write_bits(0xFF, 8); writer.write_bits(0xAA, 8); let data = writer.finish();
assert_eq!(data, vec![0xFF, 0xAA]);
}
#[test]
fn test_bit_reader_basic() {
let data = vec![0b1010_0000];
let mut reader = BitReader::new(&data);
let value = reader.read_bits(4).unwrap();
assert_eq!(value, 0b1010);
}
#[test]
fn test_bit_reader_multiple_bytes() {
let data = vec![0xFF, 0xAA];
let mut reader = BitReader::new(&data);
let first = reader.read_bits(8).unwrap();
let second = reader.read_bits(8).unwrap();
assert_eq!(first, 0xFF);
assert_eq!(second, 0xAA);
}
#[test]
fn test_roundtrip() {
let mut writer = BitWriter::new();
writer.write_bits(0b10101010, 8);
writer.write_bits(0b11001100, 8);
let data = writer.finish();
assert_eq!(data, vec![0b10101010, 0b11001100]);
let mut reader = BitReader::new(&data);
let first = reader.read_bits(8).unwrap();
let second = reader.read_bits(8).unwrap();
assert_eq!(first, 0b10101010);
assert_eq!(second, 0b11001100);
}
#[test]
fn test_write_bits_to() {
let mut data = vec![0u8; 2];
BitWriter::write_bits_to(&mut data, 0, 0b1010, 4);
BitWriter::write_bits_to(&mut data, 4, 0b1100, 4);
assert_eq!(data[0], 0b1010_1100);
}
#[test]
fn test_read_bits_at() {
let data = vec![0b1010_1100];
let first_nibble = BitReader::read_bits_at(&data, 0, 4);
let second_nibble = BitReader::read_bits_at(&data, 4, 4);
assert_eq!(first_nibble, 0b1010);
assert_eq!(second_nibble, 0b1100);
}
}