use crate::fixedpoint::{add, shl, shr, sub};
pub struct BitstreamReader<'a> {
data: &'a [i16],
pos: usize,
bits_remaining: i16,
current_word: i16,
pub total_bits_remaining: i16,
pub last_bit: i16,
}
impl<'a> BitstreamReader<'a> {
pub fn new(data: &'a [i16], total_bits: i16) -> Self {
BitstreamReader {
data,
pos: 0,
bits_remaining: 0,
current_word: 0,
total_bits_remaining: total_bits,
last_bit: 0,
}
}
pub fn read_bit(&mut self) {
if self.bits_remaining == 0 {
self.current_word = self.data[self.pos];
self.pos += 1;
self.bits_remaining = 16;
}
let new_remaining = sub(self.bits_remaining, 1);
self.bits_remaining = new_remaining;
let shifted = shr(self.current_word, new_remaining);
self.last_bit = (shifted as u16 & 1) as i16;
}
pub fn read_bits(&mut self, n: i32) -> i16 {
let mut value: i16 = 0;
for _ in 0..n {
self.read_bit();
value = add(shl(value, 1), self.last_bit);
}
value
}
}
pub struct BitstreamWriter<'a> {
data: &'a mut [i16],
pos: usize,
bits_free: i16, accumulator: u16, }
impl<'a> BitstreamWriter<'a> {
pub fn new(data: &'a mut [i16]) -> Self {
BitstreamWriter {
data,
pos: 0,
bits_free: 16,
accumulator: 0,
}
}
pub fn write_bits(&mut self, value: i16, width: i16) {
let mut remaining = width;
let mask = if width >= 16 { u16::MAX } else { (1u16 << (width as u32)) - 1 };
let mut val = (value as u16) & mask;
while remaining > 0 {
if remaining >= self.bits_free {
let shift = remaining - self.bits_free;
let bf_mask = if self.bits_free >= 16 { u16::MAX } else { (1u16 << (self.bits_free as u32)) - 1 };
self.accumulator |= (val >> (shift as u32)) & bf_mask;
if self.pos >= self.data.len() {
return; }
self.data[self.pos] = self.accumulator as i16;
self.pos += 1;
remaining = sub(remaining, self.bits_free);
let rem_mask = if remaining >= 16 { u16::MAX } else { (1u16 << (remaining as u32)) - 1 };
val &= rem_mask;
self.accumulator = 0;
self.bits_free = 16;
} else {
let shift = self.bits_free - remaining;
self.accumulator |= val << (shift as u32);
self.bits_free = sub(self.bits_free, remaining);
remaining = 0;
}
}
}
pub fn flush(&mut self) {
if self.bits_free < 16 && self.pos < self.data.len() {
self.data[self.pos] = self.accumulator as i16;
self.pos += 1;
self.accumulator = 0;
self.bits_free = 16;
}
}
pub fn words_written(&self) -> usize {
self.pos
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_read_single_bits() {
let data = [0xA5C3u16 as i16];
let mut bs = BitstreamReader::new(&data, 16);
let expected = [1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1];
for (i, &exp) in expected.iter().enumerate() {
bs.read_bit();
assert_eq!(bs.last_bit, exp, "bit {} mismatch", i);
}
}
#[test]
fn test_read_across_words() {
let data = [0xFF00u16 as i16, 0x00FFi16];
let mut bs = BitstreamReader::new(&data, 32);
for i in 0..8 {
bs.read_bit();
assert_eq!(bs.last_bit, 1, "bit {} should be 1", i);
}
for i in 8..16 {
bs.read_bit();
assert_eq!(bs.last_bit, 0, "bit {} should be 0", i);
}
for i in 16..24 {
bs.read_bit();
assert_eq!(bs.last_bit, 0, "bit {} should be 0", i);
}
for i in 24..32 {
bs.read_bit();
assert_eq!(bs.last_bit, 1, "bit {} should be 1", i);
}
}
#[test]
fn test_read_bits_multi() {
let data = [0xA5C3u16 as i16];
let mut bs = BitstreamReader::new(&data, 16);
let val = bs.read_bits(4);
assert_eq!(val, 0b1010);
let val = bs.read_bits(4);
assert_eq!(val, 0b0101);
let val = bs.read_bits(5);
assert_eq!(val, 0b11000);
let val = bs.read_bits(3);
assert_eq!(val, 0b011);
}
#[test]
fn test_read_bits_5bit_value() {
let data = [0x7C00u16 as i16];
let mut bs = BitstreamReader::new(&data, 16);
let val = bs.read_bits(5);
assert_eq!(val, 15);
}
#[test]
fn test_writer_basic() {
let mut data = [0i16; 2];
{
let mut bw = BitstreamWriter::new(&mut data);
bw.write_bits(0b1010, 4);
bw.write_bits(0b0101, 4);
bw.write_bits(0b11000011, 8);
bw.flush();
}
assert_eq!(data[0], 0xA5C3u16 as i16);
}
#[test]
fn test_writer_cross_word() {
let mut data = [0i16; 2];
{
let mut bw = BitstreamWriter::new(&mut data);
bw.write_bits(0xFF, 8); bw.write_bits(0x00, 8); bw.write_bits(0x00, 8); bw.write_bits(0xFF, 8); bw.flush();
}
assert_eq!(data[0], 0xFF00u16 as i16);
assert_eq!(data[1], 0x00FFi16);
}
#[test]
fn test_writer_reader_roundtrip() {
let mut data = [0i16; 4];
{
let mut bw = BitstreamWriter::new(&mut data);
bw.write_bits(15, 5); bw.write_bits(7, 4); bw.write_bits(1, 1); bw.write_bits(0b101010, 6); bw.flush();
}
let mut bs = BitstreamReader::new(&data, 64);
assert_eq!(bs.read_bits(5), 15);
assert_eq!(bs.read_bits(4), 7);
assert_eq!(bs.read_bits(1), 1);
assert_eq!(bs.read_bits(6), 0b101010);
}
}