pub struct BitWriter {
output: Vec<u8>,
buffer: u64,
bits_in_buffer: u8,
}
impl BitWriter {
pub fn new() -> Self {
Self { output: Vec::with_capacity(65536), buffer: 0, bits_in_buffer: 0 }
}
pub fn with_capacity(capacity: usize) -> Self {
Self { output: Vec::with_capacity(capacity), buffer: 0, bits_in_buffer: 0 }
}
#[inline]
fn flush_bytes(&mut self) {
while self.bits_in_buffer >= 8 {
self.output.push(self.buffer as u8);
self.buffer >>= 8;
self.bits_in_buffer -= 8;
}
}
#[inline]
pub fn write_bits(&mut self, value: u32, n: u8) {
debug_assert!(n <= 32);
if n == 0 {
return;
}
self.buffer |= (value as u64) << self.bits_in_buffer;
self.bits_in_buffer += n;
if self.bits_in_buffer >= 32 {
self.flush_bytes();
}
}
#[inline]
pub fn write_bit(&mut self, bit: bool) {
self.write_bits(bit as u32, 1);
}
pub fn write_bits_reversed(&mut self, code: u32, length: u8) {
let reversed = reverse_bits(code, length);
self.write_bits(reversed, length);
}
pub fn align_to_byte(&mut self) {
if self.bits_in_buffer % 8 != 0 {
self.bits_in_buffer = ((self.bits_in_buffer + 7) / 8) * 8;
}
self.flush_bytes();
}
pub fn write_byte(&mut self, byte: u8) {
if self.bits_in_buffer % 8 == 0 {
self.flush_bytes();
self.output.push(byte);
} else {
self.write_bits(byte as u32, 8);
}
}
pub fn write_u16_le(&mut self, value: u16) {
self.write_byte(value as u8);
self.write_byte((value >> 8) as u8);
}
pub fn write_u32_le(&mut self, value: u32) {
self.write_byte(value as u8);
self.write_byte((value >> 8) as u8);
self.write_byte((value >> 16) as u8);
self.write_byte((value >> 24) as u8);
}
pub fn write_bytes(&mut self, bytes: &[u8]) {
for &b in bytes {
self.write_byte(b);
}
}
pub fn finish(mut self) -> Vec<u8> {
self.align_to_byte();
self.output
}
pub fn len(&self) -> usize {
self.output.len() + if self.bits_in_buffer > 0 { 1 } else { 0 }
}
pub fn is_empty(&self) -> bool {
self.output.is_empty() && self.bits_in_buffer == 0
}
pub fn as_bytes(&self) -> &[u8] {
&self.output
}
pub fn clear(&mut self) {
self.output.clear();
self.buffer = 0;
self.bits_in_buffer = 0;
}
}
impl Default for BitWriter {
fn default() -> Self {
Self::new()
}
}
pub(crate) fn reverse_bits(value: u32, n: u8) -> u32 {
let mut result = 0u32;
let mut v = value;
for _ in 0..n {
result = (result << 1) | (v & 1);
v >>= 1;
}
result
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_write_bits() {
let mut writer = BitWriter::new();
writer.write_bits(0b011, 3); writer.write_bits(0b11010, 5); let output = writer.finish();
assert_eq!(output, vec![0xD3]); }
#[test]
fn test_write_cross_byte() {
let mut writer = BitWriter::new();
writer.write_bits(0xFFF, 12); let output = writer.finish();
assert_eq!(output, vec![0xFF, 0x0F]);
}
#[test]
fn test_write_u16_le() {
let mut writer = BitWriter::new();
writer.write_u16_le(0x1234);
let output = writer.finish();
assert_eq!(output, vec![0x34, 0x12]);
}
#[test]
fn test_reverse_bits() {
assert_eq!(reverse_bits(0b1100, 4), 0b0011);
assert_eq!(reverse_bits(0b10101, 5), 0b10101);
assert_eq!(reverse_bits(0b11110000, 8), 0b00001111);
}
#[test]
fn test_write_bits_reversed() {
let mut writer = BitWriter::new();
writer.write_bits_reversed(0b1100, 4);
writer.write_bits(0, 4); let output = writer.finish();
assert_eq!(output[0] & 0x0F, 0b0011);
}
}