huffman-compression-rust/src/huffman.rs

use std::{
    cmp::Reverse,
    collections::{BinaryHeap, HashMap},
};
use Tree::*;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Tree<T> {
    Leaf {
        freq: u64,
        token: T,
    },
    Node {
        freq: u64,
        left: Box<Tree<T>>,
        right: Box<Tree<T>>,
    },
}

#[allow(dead_code)]
impl<T: Clone> Tree<T> {
    pub fn freq(&self) -> u64 {
        match self {
            Leaf { freq, .. } => *freq,
            Node { freq, .. } => *freq,
        }
    }

    pub fn token(&self) -> Option<T> {
        match self {
            Leaf { token, .. } => Some(token.clone()),
            Node { .. } => None,
        }
    }

    pub fn left(&self) -> Option<&Tree<T>> {
        match self {
            Node { left, .. } => Some(left),
            Leaf { .. } => None,
        }
    }

    pub fn right(&self) -> Option<&Tree<T>> {
        match self {
            Node { right, .. } => Some(right),
            Leaf { .. } => None,
        }
    }
}

impl<T: Clone + Eq> Ord for Tree<T> {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.freq().cmp(&other.freq())
    }
}

impl<T: Clone + Eq> PartialOrd for Tree<T> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

pub fn huffman_tree<T: Eq + Clone>(freqs: &HashMap<T, u64>) -> Tree<T> {
    let mut heap = BinaryHeap::new();
    for (token, freq) in freqs {
        let (freq, token) = (*freq, token.clone());
        heap.push(Reverse(Leaf { freq, token }))
    }

    while heap.len() > 1 {
        let node1 = heap.pop().unwrap().0;
        let node2 = heap.pop().unwrap().0;

        let merged_node = Node {
            freq: node1.freq() + node2.freq(),
            left: Box::new(node1),
            right: Box::new(node2),
        };
        heap.push(Reverse(merged_node));
    }

    heap.pop().unwrap().0
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::freq::char_frequencies;

    #[test]
    fn learn_frequencies_test() {
        let input = vec!["this is an epic sentence".to_string(), "xyz ".to_string()];
        let freqs = char_frequencies(&input);
        assert_eq!(freqs[&' '], 5);
        assert_eq!(freqs[&'t'], 2);
        assert_eq!(freqs[&'i'], 3);
        assert_eq!(freqs[&'p'], 1);
        assert_eq!(freqs[&'z'], 1);
        assert_eq!(freqs.keys().len(), 13);
    }

    #[test]
    fn huffman_tree_test() {
        let mut freqs = HashMap::new();
        freqs.insert('a', 40);
        freqs.insert('b', 35);
        freqs.insert('c', 20);
        freqs.insert('d', 5);

        let tree = huffman_tree(&freqs);
        assert_eq!(tree.freq(), 100);

        // the most frequent character only requires 1 bit
        assert_eq!(tree.left().and_then(|n| n.token()), Some('a'));
        assert_eq!(tree.left().map(|n| n.freq()), Some(40));

        // the second most frequent character requires 2 bits
        assert_eq!(
            tree.right().and_then(|t| t.right()).and_then(|n| n.token()),
            Some('b')
        );
        assert_eq!(
            tree.right().and_then(|t| t.right()).map(|n| n.freq()),
            Some(35)
        );

        // the least frequent characters require 3 bits
        assert_eq!(
            tree.right()
                .and_then(|t| t.left())
                .and_then(|t| t.left())
                .and_then(|n| n.token()),
            Some('d')
        );
        assert_eq!(
            tree.right()
                .and_then(|t| t.left())
                .and_then(|t| t.left())
                .map(|n| n.freq()),
            Some(5)
        );

        assert_eq!(
            tree.right()
                .and_then(|t| t.left())
                .and_then(|t| t.right())
                .and_then(|n| n.token()),
            Some('c')
        );
        assert_eq!(
            tree.right()
                .and_then(|t| t.left())
                .and_then(|t| t.right())
                .map(|n| n.freq()),
            Some(20)
        );
    }
}