use super::read_byte::ReadByte as _;
use std::{cmp::Ordering, io};
const SLACK: usize = 6;
#[derive(Copy, Clone)]
pub(crate) struct Node {
pub(crate) byte: u8,
pub(crate) flag: i8,
pub(crate) one: isize,
pub(crate) zero: isize,
}
impl Default for Node {
fn default() -> Self {
Self {
byte: Default::default(),
flag: -1,
one: -1,
zero: -1,
}
}
}
pub(crate) struct HuffmanTree<const T: usize> {
pub(crate) entries: Vec<Node>,
}
impl<const T: usize> HuffmanTree<T> {}
#[derive(Default, Copy, Clone, Debug)]
struct TreeEntry {
value: i32,
bit_length: usize,
}
impl<const T: usize> HuffmanTree<T> {
pub const TREE_SIZE: usize = 2 * T;
pub fn new() -> Self {
Self {
entries: vec![Default::default(); Self::TREE_SIZE],
}
}
pub fn read_from<R: io::Read>(&mut self, reader: &mut R) -> io::Result<()> {
let tree_bytes: usize = reader.read_byte()? as usize;
assert!(T >= tree_bytes * 2);
let mut tree_count = [0i32; 32];
let mut i = 0;
let mut tree_max_length = 0;
let mut tree_entries = 0;
let mut tree_entry: [TreeEntry; T] = [Default::default(); T];
let mut hufftree: Vec<Node> = vec![Default::default(); 2 * T + SLACK];
for _ in (0..tree_bytes).rev() {
let current = reader.read_byte()?;
let first_node = current >> 4;
let second_node = current & 0x0F;
macro_rules! insert {
($value: ident) => {
if ($value != 0) {
if ($value > tree_max_length) {
tree_max_length = $value;
}
tree_count[$value as usize] += 1;
tree_entry[tree_entries].value = i;
tree_entry[tree_entries].bit_length = $value as usize;
tree_entries += 1;
}
i += 1;
};
}
insert!(first_node);
insert!(second_node);
}
let mut j = 0;
for i in 0..=tree_max_length {
j = (j << 1) + tree_count[i as usize];
}
j = (1 << tree_max_length) - j;
for _ in 0..j {
tree_entry[tree_entries].value = T as i32;
tree_entry[tree_entries].bit_length = tree_max_length as usize;
tree_entries += 1;
}
tree_entry[0..tree_entries].sort_by(|a, b| match a.bit_length.cmp(&b.bit_length) {
Ordering::Equal => a.value.cmp(&b.value),
v => v,
});
let mut i: isize = tree_entries as isize - 1;
let mut lvlstart = T * 2 + SLACK - 1;
let mut next = lvlstart;
for codelen in (1..=tree_max_length).rev() {
while i >= 0 && tree_entry[i as usize].bit_length == codelen as usize {
hufftree[next].byte = tree_entry[i as usize].value as u8;
hufftree[next].flag = 1;
next -= 1;
i -= 1;
}
let parents = next;
if codelen > 1 {
let range = (((parents + 2) as i16)..=(lvlstart as i16)).step_by(2);
for j in range.rev() {
hufftree[next].one = j.into();
hufftree[next].zero = (j - 1).into();
hufftree[next].flag = 0;
next -= 1;
}
}
lvlstart = parents;
}
hufftree[0].one = (next as isize) + 2;
hufftree[0].zero = (next as isize) + 1;
hufftree[0].flag = 0;
self.entries = hufftree;
Ok(())
}
}