# Huffman encoding in Rust

An implementation of Huffman encoding / decoding as a learning exercise. I'm still new to Rust, so any tips about making things more natural for the language are appreciated.

mod huff {

use std::boxed::Box;
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::collections::BinaryHeap;
use std::fmt::Write;

pub struct SymbolFreq {
symbol: char,
frequency: u32,
}

fn make_frequency_table(input: &String) -> Vec<SymbolFreq> {

let mut freqs = BTreeMap::new();

for c in input.chars() {
*freqs.entry(c).or_insert(0) += 1;
}

let mut result = freqs.iter().map(|(k, v)| SymbolFreq{ symbol: *k, frequency: *v }).collect::<Vec<_>>();

result.sort_by(|a, b| a.frequency.cmp(&b.frequency));

result
}

pub struct Node {
left : Option<Box<Node>>,
right : Option<Box<Node>>,
symbol_freq : SymbolFreq,
}

impl Node {
fn is_leaf(&self) -> bool {
self.left.is_none() && self.right.is_none()
}
}

impl PartialEq for Node {
fn eq(&self, other: &Self) -> bool {
self.symbol_freq.frequency.eq(&other.symbol_freq.frequency)
}
}

impl Eq for Node { }

impl PartialOrd for Node {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
other.symbol_freq.frequency.partial_cmp(&self.symbol_freq.frequency)
}
}

impl Ord for Node {
fn cmp(&self, other: &Self) -> Ordering {
other.symbol_freq.frequency.cmp(&self.symbol_freq.frequency)
}
}

fn make_heap_node(symbol_freq: SymbolFreq) -> Box<Node> {
Box::new(Node{ left: None, right: None, symbol_freq: symbol_freq })
}

fn make_combined_heap_node(left: Box<Node>, right: Box<Node>) -> Box<Node> {

let symbol_freq = SymbolFreq{
symbol: 0 as char,
frequency: left.symbol_freq.frequency + right.symbol_freq.frequency
};

Box::new(Node{ left: Some(left), right: Some(right), symbol_freq: symbol_freq })
}

fn make_queue(input: Vec<SymbolFreq>) -> BinaryHeap<Box<Node>> {

let mut freqs = BinaryHeap::with_capacity(input.len());

for e in input {
freqs.push(make_heap_node(e));
}

freqs
}

fn make_tree(mut input: BinaryHeap<Box<Node>>) -> Option<Box<Node>> {

if input.is_empty() {
return None;
}

let root;

loop {

assert!(!input.is_empty());

let l = input.pop().unwrap();

let r = match input.pop() {
Some(x) => x,
None => { root = l; break; },
};

let n = make_combined_heap_node(l, r);

match input.is_empty() {
true => { root = n; break; },
false => input.push(n),
}
}

assert!(input.is_empty());

Some(root)
}

fn make_encoding_table(input: &Option<Box<Node>>, mut encoding: Vec<bool>, output: &mut BTreeMap<char, Vec<bool>>) {

let n = match input.as_ref() {
Some(x) => x,
None => return,
};

if n.is_leaf() {

if encoding.is_empty() { // root is leaf! -> add extra bit
encoding.push(false);
}

assert!(output.insert(n.symbol_freq.symbol, encoding).is_none());
}
else {

let mut l = encoding.clone();
l.push(false);
make_encoding_table(&n.left, l, output);

let mut r = encoding;
r.push(true);
make_encoding_table(&n.right, r, output);
}
}

fn make_encodings(input: &String) -> (Option<Box<Node>>, BTreeMap<char, Vec<bool>>) {

let freqs = make_frequency_table(input);
debug::print_frequencies(&freqs);

let queue = make_queue(freqs);
debug::print_queue(&queue);

let root = make_tree(queue);
debug::print_tree(&root);

let mut encodings = BTreeMap::new();
make_encoding_table(&root, Vec::new(), &mut encodings);

(root, encodings)
}

fn do_encoding(input: &String, encodings: &BTreeMap<char, Vec<bool>>) -> Vec<bool> {

let mut result = Vec::new();

for c in input.chars() {
result.extend(encodings.get(&c).unwrap());
}

result
}

fn do_decoding(input: &Vec<bool>, root: &Option<Box<Node>>) -> Result<String, &'static str> {

let mut result = String::new();
let mut i = 0usize;

while i != input.len() {

let mut next = root;

loop {

let node = match next.as_ref() {
Some(x) => x,
None => return Err("do_decoding() failed: reached null node in tree."),
};

if node.is_leaf() {
result.push(node.symbol_freq.symbol);
break;
}

if i == input.len() {
return Err("do_decoding() failed: reached end of input partway through decoding char.");
}

next = match input[i] {
false => &node.left,
true => &node.right,
};

i += 1;
}

if next == root { // root is leaf! -> always consume at least one bit!
i += 1;
}
}

Ok(result)
}

pub type Tree = Option<Box<Node>>;
pub type EncodedData = Vec<bool>;

pub fn encode(input: &String) -> (Tree, EncodedData) {

let (tree, encodings) = make_encodings(input);
debug::print_encodings(&encodings);

let encoded_data = do_encoding(input, &encodings);
debug::print_byte_comparison(input, &encoded_data);

(tree, encoded_data)
}

pub fn decode(input: &EncodedData, tree: &Tree) -> Result<String, &'static str> {
do_decoding(input, tree)
}

mod debug {

use super::*;
use std::collections::BTreeMap;
use std::collections::BinaryHeap;
use std::fmt::Write;

pub fn print_frequencies(frequencies: &Vec<SymbolFreq>) {
println!("frequencies:");
for e in frequencies.iter() {
println!("\t({}, {})", e.symbol, e.frequency);
}
}

pub fn print_queue(queue: &BinaryHeap<Box<Node>>) {
println!("queue:");
for n in queue {
println!("\t({}, {})", n.symbol_freq.symbol, n.symbol_freq.frequency);
}
}

fn print_node(node: &Option<Box<Node>>, indent: u32) {
let n = match node.as_ref() {
Some(x) => x,
None => return,
};

for _ in 0..indent {
print!("\t");
}
println!("({}, {})", n.symbol_freq.symbol, n.symbol_freq.frequency);

print_node(&n.left, indent + 1);
print_node(&n.right, indent + 1);
}

pub fn print_tree(root: &Option<Box<Node>>) {
println!("tree:");
print_node(&root, 0);
}

fn encoding_to_string(encoding: &Vec<bool>) -> String {
let mut s = String::with_capacity(encoding.len());
for b in encoding {
write!(&mut s, "{}", *b as u8).unwrap();
}
s
}

pub fn print_encodings(encodings: &BTreeMap<char, Vec<bool>>) {
println!("encodings:");
for (k, v) in encodings {
print!("{}: {}\n", k, encoding_to_string(v));
}
}

pub fn print_byte_comparison(unencoded: &String, encoded: &Vec<bool>) {
print!("data: {}\n", encoding_to_string(&encoded));
print!("original: {}bytes\nencoded: ~{}bytes\n", unencoded.len(), encoded.len() / 8);
}

} // debug

#[cfg(test)]
mod tests {

use super::*;

#[test]
fn encode_decode_empty_input() {
let input = "".to_string();
let (t, d) = encode(&input);

assert!(t.is_none());
assert!(d.is_empty());

let o = decode(&d, &t).unwrap();

assert_eq!(o, input);
}

#[test]
fn encode_decode_single_char_input() {
let input = "aaaaaaa".to_string();
let (t, d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

let o = decode(&d, &t).unwrap();

assert_eq!(o, input);
}

#[test]
fn encode_decode_two_char_input() {
let input = "abaabbaa".to_string();
let (t, d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

let o = decode(&d, &t).unwrap();

assert_eq!(o, input);
}

#[test]
fn encode_decode_three_char_input() {
let input = "   abaa    bba a".to_string();
let (t, d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

let o = decode(&d, &t).unwrap();

assert_eq!(o, input);
}

#[test]
fn encode_decode_utf8_input() {
let input = "𝄞 αβ 忠犬ハチ公 hi".to_string();
let (t, d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

let o = decode(&d, &t).unwrap();

assert_eq!(o, input);
}

#[test]
fn decode_empty_encoded_data() {

let input = "𝄞 αβ 忠犬ハチ公 hi asdkfweiuryhfiusdf".to_string();
let (t, d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

let o = decode(&Vec::new(), &t).unwrap();

assert!(o.is_empty());
}

#[test]
#[should_panic(expected = "reached end of input")]
fn decode_corrupt_encoded_data() {

let input = "𝄞 αβ 忠犬ハチ公 hi asdkfweiuryhfiusdf".to_string();
let (t, mut d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

d.pop();
d.pop();

let _ = decode(&d, &t).unwrap();
}

#[test]
#[should_panic(expected = "reached null node")]
fn decode_corrupt_tree() {

let input = "𝄞 αβ 忠犬ハチ公 hi asdkfweiuryhfiusdf".to_string();
let (mut t, d) = encode(&input);

assert!(!t.is_none());
assert!(!d.is_empty());

t.as_mut().unwrap().left = None;

let _ = decode(&d, &t).unwrap();
}

} // tests

} // huff

#[cfg(not(test))]
fn main() {

let input = "this is an example for huffman encoding";
print!("input: {}\n", input);

let (tree, data) = huff::encode(&input.to_string());

let output = huff::decode(&data, &tree).unwrap();
print!("output: {}\n", output);
}


• It seems like a lot of trait implementations are required to put something in a BinaryHeap... is there a simpler way?
• The compiler claims std::fmt::Write is unused, but if I comment it out, it doesn't compile.
• The version of Rust I have installed is kinda old (rustc 1.3.0 (9a92aaf19 2015-09-15)) as I've been without internet. Are there things that could be improved with a more modern version?