1use std;
2use std::cmp::Ordering;
3use std::collections::{BinaryHeap, VecDeque};
4
5use super::*;
6
7#[derive(Debug, Eq)]
8pub struct HuffmanType {
9 symbol: u8,
10 frequency: u64,
11}
12
13impl HuffmanType {
14 pub fn new(symbol: u8, frequency: u64) -> HuffmanType {
15 HuffmanType { symbol, frequency }
16 }
17}
18
19impl Ord for HuffmanType {
20 fn cmp(&self, other: &Self) -> Ordering {
21 (other.frequency, other.symbol).cmp(&(self.frequency, self.symbol))
22 }
23}
24
25impl PartialOrd for HuffmanType {
26 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
27 Some(self.cmp(other))
28 }
29}
30
31impl PartialEq for HuffmanType {
32 fn eq(&self, other: &Self) -> bool {
33 self.cmp(other) == Ordering::Equal
34 }
35}
36
37
38#[derive(Debug)]
39pub struct Node<T> {
40 pub value: T,
41 pub left: Option<Box<Node<T>>>,
42 pub right: Option<Box<Node<T>>>,
43 pub parent: * mut Node<T>,
44}
45
46impl <T> Node<T> {
47 pub fn new(value: T) -> Node<T> {
48 Node {
49 value,
50 left: Option::None,
51 right: Option::None,
52 parent: std::ptr::null_mut(),
53 }
54 }
55
56 pub fn set_left(&mut self, mut node: Box<Node<T>>) {
57 node.parent = self;
58 self.left = Option::Some(node);
59 }
60
61 pub fn set_right(&mut self, mut node: Box<Node<T>>) {
62 node.parent = self;
63 self.right = Option::Some(node);
64 }
65
66 pub fn is_leaf(&self) -> bool {
67 self.left.is_none() && self.right.is_none()
68 }
69}
70
71pub type HuffmanNode = Node<HuffmanType>;
72
73impl Ord for HuffmanNode {
74 fn cmp(&self, other: &Self) -> Ordering {
75 self.value.cmp(&other.value)
76 }
77}
78
79impl PartialOrd for HuffmanNode {
80 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
81 self.value.partial_cmp(&other.value)
82 }
83}
84
85impl PartialEq for HuffmanNode {
86 fn eq(&self, other: &Self) -> bool {
87 self.value.eq(&other.value)
88 }
89}
90
91impl Eq for HuffmanNode {}
92
93pub struct HuffmanTree {
94 pub root_node: Box<HuffmanNode>,
95}
96
97impl HuffmanTree {
98 pub fn new(freq_table: &[u64; NUM_BYTES]) -> Option<HuffmanTree> {
99 let mut priority_queue: BinaryHeap<Box<HuffmanNode>> = BinaryHeap::new();
100
101 for (symbol, &frequency) in freq_table.iter().enumerate() {
102 if frequency != 0 {
103 let node = HuffmanNode::new(HuffmanType::new(symbol as u8, frequency));
104
105 priority_queue.push(Box::new(node));
106 }
107 }
108
109 if priority_queue.len() == 0 {
110 return None;
111 }
112
113 while priority_queue.len() > 1 {
114 let node1 = priority_queue.pop().unwrap();
115 let node2 = priority_queue.pop().unwrap();
116
117 let mut new_node = HuffmanNode::new(
118 HuffmanType::new(0, node1.value.frequency + node2.value.frequency));
119
120 new_node.set_right(node1);
121 new_node.set_left(node2);
122
123 priority_queue.push(Box::new(new_node));
124 }
125
126 let root_node = priority_queue.pop().unwrap();
127
128 Some(HuffmanTree { root_node })
129 }
130
131 pub fn get_code_lengths(&self) -> Vec<(u8, u8)> {
132 let mut queue: VecDeque<(&HuffmanNode, u8)> = VecDeque::new();
134
135 queue.push_back((self.root_node.as_ref(), 0));
137
138 let mut code_lengths: Vec<(u8, u8)> = Vec::new();
140
141 while !queue.is_empty() {
143 let (node, depth) = queue.pop_front().unwrap();
144
145 if node.is_leaf() {
146 code_lengths.push((node.value.symbol, depth));
147 continue;
148 }
149
150 if let Some(ref left) = node.left {
151 queue.push_back((left.as_ref(), depth + 1));
152 }
153
154 if let Some(ref right) = node.right {
155 queue.push_back((right.as_ref(), depth + 1));
156 }
157 }
158
159 code_lengths
160 }
161}