#![allow(clippy::drop_non_drop)]
use std::fmt::{Debug, Formatter};
use bumpalo::Bump;
use ouroboros::self_referencing;
use thiserror::Error;
#[derive(Debug, Error)]
pub enum TryFromCodewordLengthsListError {
#[error("The codeword lengths list described an overspecified Huffman tree")]
OverspecifiedTree
}
#[derive(Debug, Error)]
pub enum VorbisHuffmanTreeWalkerError {
#[error("An attempt to use an underspecified region of the codebook Huffman tree was made")]
UnderspecifiedTree
}
#[self_referencing]
pub(super) struct VorbisHuffmanTree {
arena: Bump,
#[borrows(arena)]
#[not_covariant] root: VorbisHuffmanTreeNode<'this, VorbisHuffmanTreeEntry>
}
impl VorbisHuffmanTree {
pub(super) fn try_from_codeword_lengths<T: AsRef<[u8]>>(
codeword_lengths: T
) -> Result<Self, TryFromCodewordLengthsListError> {
let codeword_lengths = codeword_lengths.as_ref();
if codeword_lengths == [1] {
return Ok(VorbisHuffmanTreeBuilder {
arena: Bump::new(),
root_builder: |arena| VorbisHuffmanTreeNode {
left_child: Some(arena.alloc(VorbisHuffmanTreeNode {
entry: Some(VorbisHuffmanTreeEntry { number: 0 }),
..Default::default()
})),
right_child: Some(arena.alloc(VorbisHuffmanTreeNode {
entry: Some(VorbisHuffmanTreeEntry { number: 0 }),
..Default::default()
})),
entry: None
}
}
.build());
}
VorbisHuffmanTreeTryBuilder {
arena: Bump::new(),
root_builder: |arena| {
let mut root = VorbisHuffmanTreeNode::default();
for (entry_number, codeword_length) in codeword_lengths.iter().copied().enumerate()
{
if codeword_length == 0 {
continue;
}
let (entry, _) = root
.leftmost_free_leaf_at_depth(codeword_length, arena)
.ok_or(TryFromCodewordLengthsListError::OverspecifiedTree)?;
entry.entry = Some(VorbisHuffmanTreeEntry {
number: entry_number as u32
});
}
Ok(root)
}
}
.try_build()
}
pub(super) fn try_codewords_from_codeword_lengths<T: AsRef<[u64]>>(
codeword_lengths: T
) -> Result<Vec<Option<(u32, u8)>>, TryFromCodewordLengthsListError> {
let codeword_lengths = codeword_lengths.as_ref();
let mut codewords = vec![None; codeword_lengths.len()];
let mut root = VorbisHuffmanTreeNode::default();
let arena = Bump::new();
for (entry_number, codeword_length) in codeword_lengths.iter().copied().enumerate() {
let codeword_length = codeword_length as u8;
if codeword_length == 0 {
continue;
}
let (entry, codeword) = root
.leftmost_free_leaf_at_depth(codeword_length, &arena)
.ok_or(TryFromCodewordLengthsListError::OverspecifiedTree)?;
entry.entry = Some(VorbisHuffmanTreeEntry {
number: entry_number as u32
});
codewords[entry_number] = Some((codeword, codeword_length));
}
Ok(codewords)
}
pub(super) fn with_walker<R>(
&self,
f: impl FnOnce(VorbisHuffmanTreeWalker<'_, '_, VorbisHuffmanTreeEntry>) -> R
) -> R {
self.with_root(|root| f(VorbisHuffmanTreeWalker { current_node: root }))
}
}
impl Debug for VorbisHuffmanTree {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
self.with_root(|root| Debug::fmt(root, f))
}
}
#[derive(Debug)]
#[repr(transparent)]
pub(super) struct VorbisHuffmanTreeEntry {
pub(super) number: u32
}
#[derive(Debug)]
pub(super) struct VorbisHuffmanTreeWalker<'this, 'tree, V> {
current_node: &'this VorbisHuffmanTreeNode<'tree, V>
}
impl<V> VorbisHuffmanTreeWalker<'_, '_, V> {
pub(super) fn walk(
&mut self,
branch_right: bool
) -> Result<Option<&V>, VorbisHuffmanTreeWalkerError> {
self.current_node = if branch_right {
self.current_node.right_child.as_ref()
} else {
self.current_node.left_child.as_ref()
}
.ok_or(VorbisHuffmanTreeWalkerError::UnderspecifiedTree)?;
Ok(self.current_node.entry.as_ref())
}
}
#[derive(Debug)]
struct VorbisHuffmanTreeNode<'tree, V> {
left_child: Option<&'tree mut VorbisHuffmanTreeNode<'tree, V>>,
right_child: Option<&'tree mut VorbisHuffmanTreeNode<'tree, V>>,
entry: Option<V>
}
impl<'tree, V> VorbisHuffmanTreeNode<'tree, V> {
fn leftmost_free_leaf_at_depth<'this>(
&'this mut self,
depth: u8,
arena: &'tree Bump
) -> Option<(&'this mut VorbisHuffmanTreeNode<'tree, V>, u32)> {
self.leftmost_free_leaf_at_depth_internal(depth, 0, arena)
.map(|(leaf, codeword)| {
let offset = u32::BITS - depth as u32;
(leaf, codeword.wrapping_shl(offset).reverse_bits())
})
}
fn leftmost_free_leaf_at_depth_internal<'this>(
&'this mut self,
depth: u8,
codeword_so_far: u32,
arena: &'tree Bump
) -> Option<(&'this mut VorbisHuffmanTreeNode<'tree, V>, u32)> {
if self.entry.is_some() {
return None;
}
if depth == 0 {
return if self.left_child.is_none() && self.right_child.is_none() {
Some((self, codeword_so_far))
} else {
None
};
}
self.left_child
.get_or_insert_with(|| arena.alloc(Default::default()))
.leftmost_free_leaf_at_depth_internal(depth - 1, codeword_so_far, arena)
.or_else(|| {
self.right_child
.get_or_insert_with(|| arena.alloc(Default::default()))
.leftmost_free_leaf_at_depth_internal(
depth - 1,
codeword_so_far | 1 << (depth - 1),
arena
)
})
}
}
impl<V> Default for VorbisHuffmanTreeNode<'_, V> {
fn default() -> Self {
Self {
left_child: None,
right_child: None,
entry: None
}
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn huffman_tree_from_codeword_lengths_works() {
let tree = VorbisHuffmanTree::try_from_codeword_lengths(&[2, 4, 4, 4, 4, 2, 3, 3])
.expect("The Huffman tree was assumed to not be overspecified");
tree.with_root(|root| eprintln!("Tree: {:#?}", root));
eprintln!(
"Tree nodes arena allocated bytes: {}",
tree.borrow_arena().allocated_bytes()
);
for (entry_number, (codeword, codeword_length)) in [
(0b00, 2),
(0b0100, 4),
(0b0101, 4),
(0b0110, 4),
(0b0111, 4),
(0b10, 2),
(0b110, 3),
(0b111, 3)
]
.into_iter()
.enumerate()
{
eprintln!("Testing decode of codeword {:0codeword_length$b}", codeword);
tree.with_walker(|mut walker| {
let mut read_entry = None;
for i in (0..codeword_length).rev() {
let bit = codeword >> i & 1;
read_entry = walker
.walk(bit != 0)
.expect("The Huffman tree was assumed to not be underspecified");
}
assert_eq!(
read_entry
.expect("The Huffman tree could not decode an assumed valid codeword")
.number,
entry_number as u32
);
});
}
}
#[test]
fn single_entry_huffman_tree_works() {
let tree = VorbisHuffmanTree::try_from_codeword_lengths(&[1])
.expect("The Huffman tree was assumed to not be overspecified");
tree.with_root(|root| eprintln!("Tree: {:#?}", root));
eprintln!(
"Tree nodes arena allocated bytes: {}",
tree.borrow_arena().allocated_bytes()
);
for bit in [false, true] {
assert_eq!(
tree.with_walker(|mut walker| walker
.walk(bit)
.expect("The Huffman tree was assumed to not be underspecified")
.expect("A single-bit codeword was expected to be decoded")
.number),
0
);
}
}
#[test]
fn overspecified_huffman_tree_is_rejected() {
VorbisHuffmanTree::try_from_codeword_lengths(&[2, 4, 4, 4, 4, 2, 3, 3, 32])
.expect_err("The Huffman tree was assumed to be overspecified");
}
#[test]
fn underspecified_huffman_tree_codewords_are_rejected() {
let tree = VorbisHuffmanTree::try_from_codeword_lengths(&[2, 4, 4, 4, 2, 3, 3])
.expect("The Huffman tree was assumed to not be overspecified");
tree.with_root(|root| eprintln!("Tree: {:#?}", root));
eprintln!(
"Tree nodes arena allocated bytes: {}",
tree.borrow_arena().allocated_bytes()
);
for codeword in [
&[false, true, true, true][..],
&[false, true, false, false, false][..]
] {
tree.with_walker(|mut walker| {
for codeword_bit in codeword.iter().take(codeword.len() - 1) {
walker
.walk(*codeword_bit)
.expect("Unexpected underspecified Huffman tree error");
}
walker
.walk(*codeword.last().unwrap())
.expect_err("Expected underspecified Huffman tree error");
});
}
}
#[test]
fn empty_huffman_tree_is_underspecified() {
let tree = VorbisHuffmanTreeBuilder {
arena: Bump::new(),
root_builder: |_| Default::default()
}
.build();
tree.with_walker(|mut walker| {
walker
.walk(false) .expect_err("Expected underspecified Huffman tree error")
});
}
#[test]
#[ignore = "Takes a long time to run"]
fn monstrous_codeword_lengths_list_has_reasonable_resource_consumption() {
const MONSTROUS_CODEWORD_LENGTH: u8 = 16;
let tree = VorbisHuffmanTree::try_from_codeword_lengths(
&[MONSTROUS_CODEWORD_LENGTH; 2_usize.pow(MONSTROUS_CODEWORD_LENGTH as u32)]
)
.expect("The Huffman tree was assumed to not be overspecified");
let allocated_bytes = tree.borrow_arena().allocated_bytes();
eprintln!("Tree nodes arena allocated bytes: {}", allocated_bytes);
if allocated_bytes > 8 * 1024 * 1024 {
panic!("More than 8 MiB of RAM were allocated for the Huffman tree");
}
}
#[test]
fn codeword_lengths_list_is_assigned_expected_codewords() {
let codewords =
VorbisHuffmanTree::try_codewords_from_codeword_lengths(&[2, 4, 4, 4, 0, 4, 2, 3, 3])
.expect("The Huffman tree was assumed to not be overspecified");
const EXPECTED_CODEWORDS: &[Option<(u32, u8)>] = &[
Some((0b00, 2)),
Some((0b0010, 4)),
Some((0b1010, 4)),
Some((0b0110, 4)),
None,
Some((0b1110, 4)),
Some((0b01, 2)),
Some((0b011, 3)),
Some((0b111, 3))
];
assert_eq!(
codewords, EXPECTED_CODEWORDS,
"Unexpected codeword assignments for codeword lengths"
);
}
}