Skip to main content

huffman_rust/
huffman.rs

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        // Queue for breadth-first-search with depth
133        let mut queue: VecDeque<(&HuffmanNode, u8)> = VecDeque::new();
134
135        // Push the root node onto the queue
136        queue.push_back((self.root_node.as_ref(), 0));
137
138        // Raw code lengths
139        let mut code_lengths: Vec<(u8, u8)> = Vec::new();
140
141        // Do a breadth first search, keeping track of depth
142        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}