#![forbid(unsafe_code)]
use super::bitpack::BitReader;
use crate::AudioError;
#[derive(Debug, Clone, Default)]
pub struct CodebookEntry {
pub length: u8,
pub codeword: u32,
pub used: bool,
pub values: Vec<f32>,
}
impl CodebookEntry {
#[must_use]
pub fn unused() -> Self {
Self {
length: 0,
codeword: 0,
used: false,
values: Vec::new(),
}
}
#[must_use]
pub fn new(length: u8) -> Self {
Self {
length,
codeword: 0,
used: length > 0,
values: Vec::new(),
}
}
#[must_use]
pub fn is_valid(&self) -> bool {
self.used && self.length > 0
}
}
#[derive(Debug, Clone)]
pub enum HuffmanNode {
Internal {
left: Box<HuffmanNode>,
right: Box<HuffmanNode>,
},
Leaf(usize),
Empty,
}
impl Default for HuffmanNode {
fn default() -> Self {
Self::Empty
}
}
#[derive(Debug, Clone, Default)]
pub struct HuffmanTree {
root: HuffmanNode,
max_length: u8,
entry_count: usize,
}
impl HuffmanTree {
#[must_use]
pub fn new() -> Self {
Self {
root: HuffmanNode::Empty,
max_length: 0,
entry_count: 0,
}
}
pub fn build(lengths: &[u8]) -> Result<Self, AudioError> {
if lengths.is_empty() {
return Ok(Self::new());
}
let max_length = *lengths.iter().max().unwrap_or(&0);
if max_length > 32 {
return Err(AudioError::InvalidData("Code length too long".into()));
}
if max_length == 0 {
return Ok(Self {
root: HuffmanNode::Empty,
max_length: 0,
entry_count: lengths.len(),
});
}
let kraft_sum: u64 = lengths
.iter()
.filter(|&&l| l > 0)
.map(|&l| 1u64 << (max_length - l))
.sum();
if kraft_sum > (1u64 << max_length) {
return Err(AudioError::InvalidData(
"Invalid Huffman code lengths (Kraft inequality violated)".into(),
));
}
let entry_count = lengths.len();
let mut entries: Vec<(u8, usize)> = lengths
.iter()
.enumerate()
.filter(|(_, &l)| l > 0)
.map(|(i, &l)| (l, i))
.collect();
entries.sort_unstable();
let mut root = HuffmanNode::Empty;
let mut current_code: u32 = 0;
let mut current_length: u8 = 0;
for (len, entry_idx) in &entries {
let len = *len;
let entry_idx = *entry_idx;
if len > current_length {
current_code <<= len - current_length;
current_length = len;
}
Self::insert(&mut root, current_code, len, entry_idx)?;
current_code += 1;
}
Ok(Self {
root,
max_length,
entry_count,
})
}
fn insert(
node: &mut HuffmanNode,
code: u32,
len: u8,
entry_idx: usize,
) -> Result<(), AudioError> {
if len == 0 {
match node {
HuffmanNode::Empty => {
*node = HuffmanNode::Leaf(entry_idx);
Ok(())
}
HuffmanNode::Leaf(_) => Err(AudioError::InvalidData(
"Huffman tree: duplicate codeword".into(),
)),
HuffmanNode::Internal { .. } => Err(AudioError::InvalidData(
"Huffman tree: codeword prefix conflict".into(),
)),
}
} else {
let bit = (code >> (len - 1)) & 1;
let rest_code = code & ((1u32 << (len - 1)) - 1);
let rest_len = len - 1;
if matches!(node, HuffmanNode::Leaf(_)) {
return Err(AudioError::InvalidData(
"Huffman tree: codeword prefix is already a leaf".into(),
));
}
if matches!(node, HuffmanNode::Empty) {
*node = HuffmanNode::Internal {
left: Box::new(HuffmanNode::Empty),
right: Box::new(HuffmanNode::Empty),
};
}
match node {
HuffmanNode::Internal { left, right } => {
if bit == 0 {
Self::insert(left, rest_code, rest_len, entry_idx)
} else {
Self::insert(right, rest_code, rest_len, entry_idx)
}
}
_ => unreachable!("node was just set to Internal"),
}
}
}
#[must_use]
pub fn max_length(&self) -> u8 {
self.max_length
}
#[must_use]
pub fn entry_count(&self) -> usize {
self.entry_count
}
pub fn decode(&self, bits: &mut BitReader<'_>) -> Result<usize, AudioError> {
let mut node = &self.root;
loop {
match node {
HuffmanNode::Leaf(idx) => return Ok(*idx),
HuffmanNode::Empty => {
return Err(AudioError::InvalidData(
"Huffman decode: reached empty node (invalid bitstream)".into(),
))
}
HuffmanNode::Internal { left, right } => {
let bit = bits.read_bit()?;
node = if bit { right } else { left };
}
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum LookupType {
#[default]
None,
Lattice,
Tessellated,
}
impl LookupType {
#[must_use]
pub fn from_value(value: u8) -> Option<Self> {
match value {
0 => Some(LookupType::None),
1 => Some(LookupType::Lattice),
2 => Some(LookupType::Tessellated),
_ => None,
}
}
#[must_use]
pub fn has_lookup(self) -> bool {
self != LookupType::None
}
}
#[derive(Debug, Clone, Default)]
pub struct Codebook {
pub id: usize,
pub entries: usize,
pub dimensions: u16,
pub tree: HuffmanTree,
pub lookup_type: LookupType,
pub minimum_value: f32,
pub delta_value: f32,
pub value_bits: u8,
pub sequence_p: bool,
pub multiplicands: Vec<u32>,
pub entry_list: Vec<CodebookEntry>,
}
impl Codebook {
#[must_use]
pub fn new(id: usize) -> Self {
Self {
id,
..Default::default()
}
}
pub fn parse(id: usize, data: &[u8]) -> Result<Self, AudioError> {
if data.len() < 10 {
return Err(AudioError::InvalidData("Codebook data too short".into()));
}
Ok(Self::new(id))
}
#[must_use]
pub fn entry_count(&self) -> usize {
self.entries
}
#[must_use]
pub fn dimensions(&self) -> u16 {
self.dimensions
}
#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)]
pub fn decode_scalar(&self, bits: &mut BitReader<'_>) -> Result<f32, AudioError> {
let index = self.tree.decode(bits)?;
if self.entries > 0 && index >= self.entries {
return Err(AudioError::InvalidData(format!(
"scalar entry index out of range: {index} >= {}",
self.entries
)));
}
let multiplicand = self
.multiplicands
.get(index)
.copied()
.unwrap_or(index as u32);
Ok(self.minimum_value + (multiplicand as f32) * self.delta_value)
}
pub fn decode_vq(&self, bits: &mut BitReader<'_>) -> Result<Vec<f32>, AudioError> {
if !self.lookup_type.has_lookup() {
return Err(AudioError::InvalidData("Codebook has no VQ lookup".into()));
}
let index = self.tree.decode(bits)?;
self.lookup(index).ok_or_else(|| {
AudioError::InvalidData(format!(
"VQ lookup out of range: index {index}, entries {}",
self.entries
))
})
}
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub fn lookup(&self, index: usize) -> Option<Vec<f32>> {
if !self.lookup_type.has_lookup() || index >= self.entries {
return None;
}
match self.lookup_type {
LookupType::Lattice => {
let lookup_values = self.multiplicands.len();
if lookup_values == 0 {
return None;
}
let mut values = Vec::with_capacity(self.dimensions as usize);
let mut lookup_offset = index;
for _ in 0..self.dimensions {
let multiplicand = self.multiplicands[lookup_offset % lookup_values];
let value = self.minimum_value + (multiplicand as f32) * self.delta_value;
values.push(value);
lookup_offset /= lookup_values;
}
Some(values)
}
LookupType::Tessellated => {
let start = index * self.dimensions as usize;
let end = start + self.dimensions as usize;
if end > self.multiplicands.len() {
return None;
}
let values: Vec<f32> = self.multiplicands[start..end]
.iter()
.map(|&m| self.minimum_value + (m as f32) * self.delta_value)
.collect();
Some(values)
}
LookupType::None => None,
}
}
}
#[derive(Debug, Clone, Default)]
pub struct CodebookSet {
codebooks: Vec<Codebook>,
}
impl CodebookSet {
#[must_use]
pub fn new() -> Self {
Self {
codebooks: Vec::new(),
}
}
pub fn add(&mut self, codebook: Codebook) {
self.codebooks.push(codebook);
}
#[must_use]
pub fn get(&self, index: usize) -> Option<&Codebook> {
self.codebooks.get(index)
}
#[must_use]
pub fn len(&self) -> usize {
self.codebooks.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.codebooks.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::super::bitpack::{BitPacker, BitReader};
use super::*;
#[test]
fn test_codebook_entry_unused() {
let entry = CodebookEntry::unused();
assert!(!entry.used);
assert!(!entry.is_valid());
}
#[test]
fn test_codebook_entry_new() {
let entry = CodebookEntry::new(5);
assert!(entry.used);
assert_eq!(entry.length, 5);
assert!(entry.is_valid());
}
#[test]
fn test_codebook_entry_zero_length() {
let entry = CodebookEntry::new(0);
assert!(!entry.used);
assert!(!entry.is_valid());
}
#[test]
fn test_huffman_tree_new() {
let tree = HuffmanTree::new();
assert_eq!(tree.max_length(), 0);
assert_eq!(tree.entry_count(), 0);
}
#[test]
fn test_huffman_tree_build() {
let lengths = vec![2, 2, 3, 3, 3, 3];
let tree = HuffmanTree::build(&lengths).expect("should succeed");
assert_eq!(tree.max_length(), 3);
assert_eq!(tree.entry_count(), 6);
}
#[test]
fn test_huffman_tree_empty() {
let tree = HuffmanTree::build(&[]).expect("should succeed");
assert_eq!(tree.max_length(), 0);
assert_eq!(tree.entry_count(), 0);
}
#[test]
fn test_huffman_decode_basic() {
let lengths = vec![1u8, 2, 3, 3];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let mut packer = BitPacker::new();
packer.write_bits(0b0, 1);
packer.write_bits(0b01, 2);
packer.write_bits(0b011, 3);
let bytes = packer.finish();
let mut reader = BitReader::new(&bytes);
assert_eq!(tree.decode(&mut reader).expect("entry 0"), 0);
assert_eq!(tree.decode(&mut reader).expect("entry 1"), 1);
assert_eq!(tree.decode(&mut reader).expect("entry 2"), 2);
}
#[test]
fn test_huffman_decode_eof_error() {
let lengths = vec![3u8];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let data = [0x00u8];
let mut reader = BitReader::new(&data);
let _ = reader.read_bits(6); assert!(tree.decode(&mut reader).is_err());
}
#[test]
fn test_huffman_decode_single_entry() {
let lengths = vec![1u8];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let mut packer = BitPacker::new();
packer.write_bits(0b0, 1); let bytes = packer.finish();
let mut reader = BitReader::new(&bytes);
assert_eq!(tree.decode(&mut reader).expect("entry 0"), 0);
}
#[test]
fn test_lookup_type() {
assert_eq!(LookupType::from_value(0), Some(LookupType::None));
assert_eq!(LookupType::from_value(1), Some(LookupType::Lattice));
assert_eq!(LookupType::from_value(2), Some(LookupType::Tessellated));
assert_eq!(LookupType::from_value(3), None);
}
#[test]
fn test_lookup_type_has_lookup() {
assert!(!LookupType::None.has_lookup());
assert!(LookupType::Lattice.has_lookup());
assert!(LookupType::Tessellated.has_lookup());
}
#[test]
fn test_codebook_new() {
let codebook = Codebook::new(0);
assert_eq!(codebook.id, 0);
assert_eq!(codebook.entry_count(), 0);
}
#[test]
fn test_codebook_set() {
let mut set = CodebookSet::new();
assert!(set.is_empty());
set.add(Codebook::new(0));
set.add(Codebook::new(1));
assert_eq!(set.len(), 2);
assert!(!set.is_empty());
assert!(set.get(0).is_some());
assert!(set.get(1).is_some());
assert!(set.get(2).is_none());
}
#[test]
fn test_codebook_lookup_no_vq() {
let codebook = Codebook::new(0);
assert!(codebook.lookup(0).is_none());
}
#[test]
fn test_decode_scalar_with_multiplicands() {
let lengths = vec![1u8, 1];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let mut cb = Codebook::new(0);
cb.entries = 2;
cb.tree = tree;
cb.lookup_type = LookupType::None;
cb.minimum_value = 10.0;
cb.delta_value = 5.0;
cb.multiplicands = vec![0, 1];
let mut packer = BitPacker::new();
packer.write_bits(0b0, 1);
packer.write_bits(0b1, 1);
let bytes = packer.finish();
let mut reader = BitReader::new(&bytes);
let v0 = cb.decode_scalar(&mut reader).expect("entry 0");
assert!(
(v0 - 10.0_f32).abs() < 1e-6,
"entry 0 expected 10.0, got {v0}"
);
let v1 = cb.decode_scalar(&mut reader).expect("entry 1");
assert!(
(v1 - 15.0_f32).abs() < 1e-6,
"entry 1 expected 15.0, got {v1}"
);
}
#[test]
fn test_decode_scalar_no_multiplicands_uses_index() {
let lengths = vec![1u8, 1];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let mut cb = Codebook::new(0);
cb.entries = 2;
cb.tree = tree;
cb.lookup_type = LookupType::None;
cb.minimum_value = 2.0;
cb.delta_value = 3.0;
let mut packer = BitPacker::new();
packer.write_bits(0b0, 1); packer.write_bits(0b1, 1); let bytes = packer.finish();
let mut reader = BitReader::new(&bytes);
let v0 = cb.decode_scalar(&mut reader).expect("entry 0");
assert!(
(v0 - 2.0_f32).abs() < 1e-6,
"entry 0: expected 2.0, got {v0}"
);
let v1 = cb.decode_scalar(&mut reader).expect("entry 1");
assert!(
(v1 - 5.0_f32).abs() < 1e-6,
"entry 1: expected 5.0, got {v1}"
);
}
#[test]
fn test_decode_vq_lattice() {
let lengths = vec![2u8, 2, 2, 2];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let mut cb = Codebook::new(0);
cb.entries = 4;
cb.dimensions = 2;
cb.tree = tree;
cb.lookup_type = LookupType::Lattice;
cb.minimum_value = 0.0;
cb.delta_value = 1.0;
cb.multiplicands = vec![0, 1];
let mut packer = BitPacker::new();
packer.write_bits(0b01, 2); let bytes = packer.finish();
let mut reader = BitReader::new(&bytes);
let vec = cb.decode_vq(&mut reader).expect("decode vq");
assert_eq!(vec.len(), 2);
assert!(
(vec[0] - 0.0).abs() < 1e-6,
"dim 0 should be 0.0, got {}",
vec[0]
);
assert!(
(vec[1] - 1.0).abs() < 1e-6,
"dim 1 should be 1.0, got {}",
vec[1]
);
}
#[test]
fn test_decode_vq_tessellated() {
let lengths = vec![1u8, 1];
let tree = HuffmanTree::build(&lengths).expect("build ok");
let mut cb = Codebook::new(0);
cb.entries = 2;
cb.dimensions = 2;
cb.tree = tree;
cb.lookup_type = LookupType::Tessellated;
cb.minimum_value = 0.0;
cb.delta_value = 0.5;
cb.multiplicands = vec![10, 20, 30, 40];
let mut packer = BitPacker::new();
packer.write_bits(0b1, 1);
let bytes = packer.finish();
let mut reader = BitReader::new(&bytes);
let vec = cb.decode_vq(&mut reader).expect("decode vq");
assert_eq!(vec.len(), 2);
assert!((vec[0] - 15.0).abs() < 1e-6, "dim 0 got {}", vec[0]);
assert!((vec[1] - 20.0).abs() < 1e-6, "dim 1 got {}", vec[1]);
}
#[test]
fn test_decode_vq_no_lookup_error() {
let mut cb = Codebook::new(0);
cb.lookup_type = LookupType::None;
let data = [0xFFu8];
let mut reader = BitReader::new(&data);
assert!(cb.decode_vq(&mut reader).is_err());
}
#[test]
fn test_huffman_build_kraft_violation() {
let lengths = vec![1u8, 1, 1];
assert!(HuffmanTree::build(&lengths).is_err());
}
}