rpeg/
huffman_tree.rs

1pub struct HuffmanTree {
2    nodes: Vec<Node>,
3
4    pub table_type: u8, // 0 for luminance, 1 for chrominance
5    pub table_number: u8,
6
7    depth: usize,
8}
9
10type NodeIndex = usize;
11
12impl HuffmanTree {
13    pub fn new(table_type: u8, table_number: u8) -> Self {
14        Self {
15            nodes: Vec::new(),
16            table_type,
17            table_number,
18            depth: 0,
19        }
20    }
21
22    pub fn build(&mut self, lengths: &Vec<u8>, vals: &Vec<u8>) {
23        self.nodes = Vec::new();
24
25        self.depth = lengths.len();
26
27        let root = Node::new(None, None, None, 0);
28        self.nodes.push(root);
29
30        self.insert_children(0);
31
32        let mut leftmost_node_index: NodeIndex = 1;
33
34        let mut current_value_index = 0;
35
36        for l in lengths {
37            let mut current_node_index = leftmost_node_index;
38            for _ in 0..*l {
39                self.nodes[current_node_index].code = vals[current_value_index];
40                current_value_index += 1;
41                current_node_index = self
42                    .right_node_level(current_node_index)
43                    .expect("[E] - length exceeded current's level space");
44            }
45
46            let leftmost_parent_index = current_node_index;
47
48            while let Some(next_node_index) = self.right_node_level(current_node_index) {
49                self.insert_children(current_node_index);
50                current_node_index = next_node_index;
51            }
52            self.insert_children(current_node_index);
53
54            leftmost_node_index = self.nodes[leftmost_parent_index].left_child.unwrap();
55        }
56    }
57
58    fn insert_children(&mut self, parent_node: NodeIndex) {
59        let left_child = Node::new(Some(parent_node), None, None, 0);
60        self.nodes.push(left_child);
61        self.nodes[parent_node].left_child = Some(self.nodes.len() - 1);
62
63        let right_child = Node::new(Some(parent_node), None, None, 0);
64        self.nodes.push(right_child);
65        self.nodes[parent_node].right_child = Some(self.nodes.len() - 1);
66    }
67
68    fn right_node_level(&self, node: NodeIndex) -> Option<NodeIndex> {
69        let parent = self.nodes[node].parent_index;
70
71        let parent = match parent {
72            Some(parent) => parent,
73            None => return None,
74        };
75
76        let right_sibling = self.nodes[parent].right_child.unwrap();
77
78        if right_sibling == node {
79            match self.right_node_level(parent) {
80                Some(parent_right_sibling) => {
81                    Some(self.nodes[parent_right_sibling].left_child.unwrap())
82                }
83                None => None,
84            }
85        } else {
86            Some(right_sibling)
87        }
88    }
89
90    pub fn try_decode(&self, bits: &[bool]) -> HuffmanResult {
91        let mut node_index = 0;
92
93        for bit in bits.iter() {
94            node_index = if *bit {
95                self.nodes[node_index].right_child.unwrap()
96            } else {
97                self.nodes[node_index].left_child.unwrap()
98            };
99        }
100
101        if self.nodes[node_index].is_leaf() {
102            if bits.len() > self.depth || self.nodes[node_index].code == 0 {
103                return HuffmanResult::EOB;
104            }
105            return HuffmanResult::Some(self.nodes[node_index].code);
106        } else {
107            return HuffmanResult::None;
108        }
109    }
110
111    pub fn print(&self) {
112        let mut stack = Vec::new();
113
114        stack.push((0, String::from("")));
115        println!("depth: {}", self.depth);
116
117        while let Some((node, node_code)) = stack.pop() {
118            if self.nodes[node].is_leaf() {
119                println!("{:02X} = {}", self.nodes[node].code, node_code);
120            }
121            if let Some(left_node) = self.nodes[node].left_child {
122                stack.push((left_node, format!("{}0", node_code)));
123            }
124            if let Some(right_node) = self.nodes[node].right_child {
125                stack.push((right_node, format!("{}1", node_code)));
126            }
127        }
128    }
129}
130
131pub enum HuffmanResult {
132    Some(u8),
133    None,
134    EOB,
135}
136
137#[derive(PartialEq, Eq)]
138struct Node {
139    parent_index: Option<NodeIndex>,
140
141    left_child: Option<NodeIndex>,
142    right_child: Option<NodeIndex>,
143    code: u8,
144}
145
146impl Node {
147    fn new(
148        parent_index: Option<NodeIndex>,
149        left_child: Option<NodeIndex>,
150        right_child: Option<NodeIndex>,
151        code: u8,
152    ) -> Self {
153        Self {
154            parent_index,
155            left_child,
156            right_child,
157            code,
158        }
159    }
160
161    fn is_leaf(&self) -> bool {
162        self.left_child.is_none() && self.right_child.is_none()
163    }
164}
rpeg/huffman_tree.rs

rpeg/
huffman_tree.rs
