use super::tree::{HuffmanNode, HuffmanTree};
use crate::error::{Result, ZiporaError};
#[derive(Debug)]
pub(crate) struct BitStreamReader<'a> {
pub(crate) data: &'a [u8],
pub(crate) current: u64,
pub(crate) bit_count: usize,
pub(crate) byte_pos: usize,
}
impl<'a> BitStreamReader<'a> {
pub(crate) fn new(data: &'a [u8]) -> Self {
let mut reader = Self {
data,
current: 0,
bit_count: 0,
byte_pos: 0,
};
reader.refill();
reader
}
#[inline]
pub(crate) fn refill(&mut self) {
while self.bit_count <= 56 && self.byte_pos < self.data.len() {
self.current |= (self.data[self.byte_pos] as u64) << self.bit_count;
self.bit_count += 8;
self.byte_pos += 1;
}
}
#[inline]
pub(crate) fn peek(&self, count: usize) -> u64 {
debug_assert!(count <= 64);
self.current & ((1u64 << count) - 1)
}
#[inline]
pub(crate) fn consume(&mut self, count: usize) {
debug_assert!(count <= self.bit_count);
self.current >>= count;
self.bit_count -= count;
}
#[inline]
#[allow(dead_code)]
pub(crate) fn read(&mut self, count: usize) -> u64 {
if count > self.bit_count {
self.refill();
}
let result = self.peek(count);
self.consume(count);
result
}
#[inline]
#[allow(dead_code)]
pub(crate) fn has_bits(&self) -> bool {
self.bit_count > 0 || self.byte_pos < self.data.len()
}
#[inline]
#[allow(dead_code)]
pub(crate) fn remaining_bits(&self) -> usize {
self.bit_count + (self.data.len() - self.byte_pos) * 8
}
}
#[derive(Debug)]
pub struct HuffmanDecoder {
tree: HuffmanTree,
}
impl HuffmanDecoder {
pub fn new(tree: HuffmanTree) -> Self {
Self { tree }
}
pub fn decode(&self, encoded_data: &[u8], output_length: usize) -> Result<Vec<u8>> {
if encoded_data.is_empty() || output_length == 0 {
return Ok(Vec::new());
}
let root = match self.tree.root() {
Some(root) => root,
None => return Err(ZiporaError::invalid_data("Empty Huffman tree")),
};
let mut result = Vec::with_capacity(output_length);
let mut current_node = root;
for &byte in encoded_data {
for bit_pos in 0..8 {
if result.len() >= output_length {
break;
}
let bit = (byte >> bit_pos) & 1 == 1;
match current_node {
HuffmanNode::Leaf { symbol, .. } => {
result.push(*symbol);
current_node = root;
current_node = match current_node {
HuffmanNode::Leaf { .. } => {
current_node
}
HuffmanNode::Internal { left, right, .. } => {
if bit {
right
} else {
left
}
}
};
}
HuffmanNode::Internal { left, right, .. } => {
current_node = if bit { right } else { left };
}
}
}
if result.len() >= output_length {
break;
}
}
if let HuffmanNode::Leaf { symbol, .. } = current_node
&& result.len() < output_length
{
result.push(*symbol);
}
if result.len() != output_length {
return Err(ZiporaError::invalid_data(format!(
"Decoded length {} != expected {}",
result.len(),
output_length
)));
}
Ok(result)
}
}