boytacean_encoding/

huffman.rs

use std::{
    cmp::Ordering,
    collections::BinaryHeap,
    io::{Cursor, Read},
    mem::size_of,
};

use boytacean_common::error::Error;

use crate::codec::Codec;

#[derive(Debug, Eq, PartialEq)]
struct Node {
    frequency: u32,
    character: Option<u8>,
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
}

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

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

pub struct Huffman;

impl Huffman {
    fn build_frequency(data: &[u8]) -> [u32; 256] {
        let mut frequency_map = [0_u32; 256];
        for &byte in data {
            frequency_map[byte as usize] += 1;
        }
        frequency_map
    }

    fn build_tree(frequency_map: &[u32; 256]) -> Option<Box<Node>> {
        let mut heap: BinaryHeap<Box<Node>> = BinaryHeap::new();

        for (byte, &frequency) in frequency_map.iter().enumerate() {
            if frequency == 0 {
                continue;
            }
            heap.push(Box::new(Node {
                frequency,
                character: Some(byte as u8),
                left: None,
                right: None,
            }));
        }

        while heap.len() > 1 {
            let left = heap.pop().unwrap();
            let right = heap.pop().unwrap();

            let merged = Box::new(Node {
                frequency: left.frequency + right.frequency,
                character: None,
                left: Some(left),
                right: Some(right),
            });

            heap.push(merged);
        }

        heap.pop()
    }

    fn build_codes(node: &Node, prefix: Vec<u8>, codes: &mut [Vec<u8>]) {
        if let Some(character) = node.character {
            codes[character as usize] = prefix;
        } else {
            if let Some(ref left) = node.left {
                let mut left_prefix = prefix.clone();
                left_prefix.push(0);
                Self::build_codes(left, left_prefix, codes);
            }
            if let Some(ref right) = node.right {
                let mut right_prefix = prefix;
                right_prefix.push(1);
                Self::build_codes(right, right_prefix, codes);
            }
        }
    }

    fn encode_data(data: &[u8], codes: &[Vec<u8>]) -> Vec<u8> {
        let mut bit_buffer = Vec::new();
        let mut current_byte = 0u8;
        let mut bit_count = 0;

        for &byte in data {
            let code = &codes[byte as usize];
            for &bit in code {
                current_byte <<= 1;
                if bit == 1 {
                    current_byte |= 1;
                }
                bit_count += 1;

                if bit_count == 8 {
                    bit_buffer.push(current_byte);
                    current_byte = 0;
                    bit_count = 0;
                }
            }
        }

        if bit_count > 0 {
            current_byte <<= 8 - bit_count;
            bit_buffer.push(current_byte);
        }

        bit_buffer
    }

    fn decode_data(encoded: &[u8], root: &Node, data_length: u64) -> Vec<u8> {
        let mut decoded = Vec::new();
        let mut current_node = root;
        let mut bit_index = 0;

        for &byte in encoded {
            if decoded.len() as u64 == data_length {
                break;
            }

            for bit_offset in (0..8).rev() {
                let bit = (byte >> bit_offset) & 1;
                current_node = if bit == 0 {
                    current_node.left.as_deref().unwrap()
                } else {
                    current_node.right.as_deref().unwrap()
                };

                if let Some(character) = current_node.character {
                    decoded.push(character);
                    current_node = root;
                }

                if decoded.len() as u64 == data_length {
                    break;
                }

                bit_index += 1;
                if bit_index == encoded.len() * 8 {
                    break;
                }
            }
        }

        decoded
    }

    fn encode_tree(node: &Node) -> Vec<u8> {
        let mut result = Vec::new();
        if let Some(character) = node.character {
            result.push(1);
            result.push(character);
        } else {
            result.push(0);
            if let Some(ref left) = node.left {
                result.extend(Self::encode_tree(left));
            }
            if let Some(ref right) = node.right {
                result.extend(Self::encode_tree(right));
            }
        }
        result
    }

    fn decode_tree(data: &mut &[u8]) -> Box<Node> {
        let mut node = Box::new(Node {
            frequency: 0,
            character: None,
            left: None,
            right: None,
        });

        if data[0] == 1 {
            node.character = Some(data[1]);
            *data = &data[2..];
        } else {
            *data = &data[1..];
            node.left = Some(Self::decode_tree(data));
            node.right = Some(Self::decode_tree(data));
        }
        node
    }
}

impl Codec for Huffman {
    type EncodeOptions = ();
    type DecodeOptions = ();

    fn encode(data: &[u8], _options: &Self::EncodeOptions) -> Result<Vec<u8>, Error> {
        let frequency_map = Self::build_frequency(data);
        let tree = Self::build_tree(&frequency_map)
            .ok_or(Error::CustomError(String::from("Failed to build tree")))?;

        let mut codes = vec![Vec::new(); 256];
        Self::build_codes(&tree, Vec::new(), &mut codes);

        let encoded_tree = Self::encode_tree(&tree);
        let encoded_data = Self::encode_data(data, &codes);
        let tree_length = encoded_tree.len() as u32;
        let data_length = data.len() as u64;

        let mut result = Vec::new();
        result.extend(tree_length.to_be_bytes());
        result.extend(encoded_tree);
        result.extend(data_length.to_be_bytes());
        result.extend(encoded_data);

        Ok(result)
    }

    fn decode(data: &[u8], _options: &Self::DecodeOptions) -> Result<Vec<u8>, Error> {
        let mut reader = Cursor::new(data);

        let mut buffer = [0x00; size_of::<u32>()];
        reader.read_exact(&mut buffer)?;
        let tree_length = u32::from_be_bytes(buffer);

        let mut buffer = vec![0; tree_length as usize];
        reader.read_exact(&mut buffer)?;
        let tree = Self::decode_tree(&mut buffer.as_slice());

        let mut buffer = [0x00; size_of::<u64>()];
        reader.read_exact(&mut buffer)?;
        let data_length = u64::from_be_bytes(buffer);

        let mut buffer =
            vec![0; data.len() - size_of::<u32>() - tree_length as usize - size_of::<u64>()];
        reader.read_exact(&mut buffer)?;

        let result = Self::decode_data(&buffer, &tree, data_length);

        Ok(result)
    }
}

pub fn encode_huffman(data: &[u8]) -> Result<Vec<u8>, Error> {
    Huffman::encode(data, &())
}

pub fn decode_huffman(data: &[u8]) -> Result<Vec<u8>, Error> {
    Huffman::decode(data, &())
}

#[cfg(test)]
mod tests {
    use super::{decode_huffman, encode_huffman};

    #[test]
    fn test_huffman_encoding() {
        let data = b"this is an example for huffman encoding, huffman encoding, huffman encoding";
        let encoded = encode_huffman(data).unwrap();
        let decoded = decode_huffman(&encoded).unwrap();
        assert_eq!(data.to_vec(), decoded);
        assert_eq!(encoded.len(), 109);
        assert_eq!(decoded.len(), 75);
    }
}