Skip to main content

chess/formats/chessbase/
huffman.rs

1use std::{boxed::Box, vec::Vec};
2
3use crate::formats::Bits;
4
5/// Huffman tree
6#[derive(Default)]
7pub struct Tree {
8    pub left: Option<Box<Tree>>,
9    pub right: Option<Box<Tree>>,
10    pub value: Option<u8>,
11}
12
13impl From<u8> for Tree {
14    fn from(value: u8) -> Self {
15        Self { left: None, right: None, value: Some(value) }
16    }
17}
18
19impl Tree {
20    /// Descend into the tree in the specified direction (right or left).
21    /// When a leaf is reached, push the value to result and reset the visitor.
22    // fn visit<'a>(mut current: &mut &'a Tree, root: &'a Tree, result: &mut Vec<u8>, right: bool) {
23    fn visit<'a>(
24        &'a self,
25        current: &mut &'a Tree,
26        result: &mut Vec<u8>,
27        right: bool,
28    ) -> Option<()> {
29        let direction = if right { current.right.as_ref()? } else { current.left.as_ref()? };
30        if let Some(byte) = direction.value {
31            result.push(byte);
32            *current = self;
33        } else {
34            *current = direction;
35        }
36        Some(())
37    }
38
39    #[must_use]
40    // None denotes failure
41    pub fn unpack(&self, bits: Bits<'_>, decompressed: usize) -> Option<Vec<u8>> {
42        let mut current = self;
43        let mut result = Vec::new();
44
45        for direction in bits {
46            if result.len() == decompressed {
47                break;
48            }
49
50            self.visit(&mut current, &mut result, direction)?;
51        }
52
53        Some(result)
54    }
55
56    pub fn get_mut_or_insert(&mut self, right: bool) -> &mut Self {
57        if right { self.right.get_or_insert_default() } else { self.left.get_or_insert_default() }
58    }
59
60    pub fn load(values: [(usize, u16); 256]) -> Self {
61        let mut tree = Tree::default();
62        for (i, (len, bits)) in values.into_iter().enumerate() {
63            if len > 0 {
64                let mut node = &mut tree;
65                let bits = (bits << (16 - len)).to_be_bytes();
66
67                for direction in Bits::from(bits.as_slice()).take(len) {
68                    node = node.get_mut_or_insert(direction);
69                }
70
71                node.value = Some(i as u8);
72            }
73        }
74        tree
75    }
76}