huffman/src/node.rs

use bit_vec::{self, BitVec};

#[derive(Debug, Eq, PartialEq)]
pub struct Node {
    left: Option<Box<Self>>,
    right: Option<Box<Self>>,
    character: Option<char>,
    frequency: i32,
}

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

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

impl Node {
    pub fn new() -> Node {
        Self {
            left: None,
            right: None,
            character: None,
            frequency: 0,
        }
    }

    pub fn new_with_character(character: char, frequency: i32) -> Node {
        Self {
            character: Some(character),
            left: None,
            right: None,
            frequency
        }
    }

    pub fn join_left(self: &mut Node, left: Box<Node>) {
        assert!(self.left.is_none(), "Tried to join a node to the left, but the branch was not empty.");
        self.frequency += left.get_frequency();
        self.left = Some(left);
    }

    pub fn join_right(self: &mut Node, right: Box<Node>) {
        assert!(self.right.is_none(), "Tried to join a node to the right, but the branch was not empty.");
        self.frequency += right.get_frequency();
        self.right = Some(right);
    }

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

    pub fn get_frequency(self: &Node) -> i32 {
        self.frequency
    }

    pub fn get_character(self: &Node) -> Result<char, &str> {
        if !self.is_leaf() {
            return Err("This node is not a leaf node.");
        }

        Ok(self.character.unwrap())
    }

    pub fn get_character_code(self: &Node, character: char) -> Result<BitVec, &str> {
        let output = BitVec::new();
        if self.is_leaf() && self.get_character().unwrap() == character {
            return Ok(output);
        }

        if let Some(n) = &self.left {
            let left = n.get_character_code(character);
            if left.is_ok() {
                let mut bit_result = left.unwrap();
                bit_result.push(false);
                return Ok(bit_result)
            }

        }

        if let Some(n) = &self.right {
            let right = n.get_character_code(character);
            if right.is_ok() {
                let mut bit_result = right.unwrap();
                bit_result.push(true);
                return Ok(bit_result)
            }
        }

        Err("Could not find character.")
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn make_new_node_with_character_works() {

        let example = Node::new_with_character('a', 5);
        assert!(example.is_leaf());
        assert_eq!(example.get_character().unwrap(), 'a');
    }

    #[test]
    fn make_small_tree_with_leaf_nodes() {

        let mut root = Node::new();
        let left = Node::new_with_character('a', 5);
        let right = Node::new_with_character('b', 10);


        root.join_left(Box::new(left));
        root.join_right(Box::new(right));

        assert!(!root.is_leaf());
        assert_eq!(root.get_frequency(), 15);
    }
}