hfmn 0.1.0

use std::{
    fmt::{Debug, Display},
    hash::Hash,
};

use num_derive::FromPrimitive;
use num_traits::FromPrimitive;

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
/// A prefix-free huffman code
///
/// This is essentially a customised bit vector, with a hard limit at
/// ```std::mem::sizeof::<usize>()``` bits. The bits are stored as a usize
/// underneath, and direct access to them is allowed in order for high-speed
/// implementations.
pub struct HfmnCode {
    code: usize,
    length: u8,
}

#[derive(Debug, FromPrimitive, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Hash)]
/// A single bit
pub enum Bit {
    Zero = 0,
    One = 1,
}

impl HfmnCode {
    #[must_use]
    /// Construct an empty Huffman code.
    ///
    /// # Examples
    /// ```
    /// use hfmn::HfmnCode;
    ///
    /// let code = HfmnCode::new();
    /// ```
    pub const fn new() -> Self {
        Self { code: 0, length: 0 }
    }

    #[must_use]
    /// Construct a Huffman code from a known code
    ///
    /// # Arguments
    /// * `code` - The underlying code to use, only the rightmost *length* bits
    ///   may be set.
    /// * `length` - The number of valid bits (LSB)
    pub const fn from(code: usize, length: u8) -> Self {
        Self { code, length }
    }

    /// Push a new bit into the huffman code
    ///
    /// # Panics
    /// The code must have room (i.e. ```self.len() <
    /// std::mem::size_of::<usize>() * 8```), else this will panic if asserts
    /// are enabled.
    ///
    /// # Arguments
    /// * `bit` - The bit to push into the code
    // I don't know of any platforms with a 256 byte usize, so the cast should be fine
    #[allow(clippy::cast_possible_truncation)]
    pub fn push(&mut self, bit: Bit) {
        assert!(self.length < { (std::mem::size_of::<usize>() * 8) as u8 });

        self.code <<= 1;
        self.code |= bit as usize;

        self.length = self.length.saturating_add(1);
    }

    #[allow(clippy::len_without_is_empty)]
    #[must_use]
    /// The current length of the code
    pub const fn len(&self) -> u8 {
        self.length
    }

    #[allow(clippy::cast_possible_truncation)]
    #[must_use]
    /// Returns a code of just the first n LSBs of the code
    ///
    /// # Panics
    /// If ```n > std::mem::size_of::<usize>() * 8```, since the code can only
    /// store that many bits
    ///
    /// # Arguments
    /// * `n` - The number of bits to copy to the new code
    pub fn first_n_bits(&self, n: usize) -> Self {
        assert!(n <= std::mem::size_of::<usize>() * 8);

        let mask = (1_usize << n).wrapping_sub(1);
        Self {
            code: self.code & mask,
            length: n as u8,
        }
    }

    #[must_use]
    /// Mutable access to the inner code
    ///
    /// Ensure that length is updated to match any relevant change in the code.
    /// The code must be zeroed beyond after its length.
    pub fn inner_mut(&mut self) -> &mut usize {
        &mut self.code
    }

    #[must_use]
    /// Mutable access to the inner length
    pub fn length_mut(&mut self) -> &mut u8 {
        &mut self.length
    }

    #[must_use]
    /// Shared access to the inner code
    pub const fn inner(&self) -> &usize {
        &self.code
    }

    #[allow(clippy::cast_possible_truncation)]
    /// Resize the code to contain exactly *size* bits
    ///
    /// Any new bits are on the LSB side
    ///
    /// # Panics
    /// If ```size > std::mem::size_of::<usize>()``` and asserts are enabled.
    ///
    /// # Arguments
    /// * `size` - The number of bits in the resized code
    pub fn resize(&mut self, size: usize) {
        assert!(size <= std::mem::size_of::<usize>() * 8);

        self.code <<= size.saturating_sub(self.length as usize);
        self.length = size as u8;
    }

    #[must_use]
    /// Get the code as an explicit array of bits
    ///
    /// # Panics
    /// If the internal code is invalid
    pub fn bits(&self) -> Vec<Bit> {
        let mut bits = Vec::new();

        for i in (0..self.length).rev() {
            let mask = 1 << i;
            bits.push(Bit::from_usize((self.code & mask) >> i).expect("Code invalid"));
        }

        bits
    }
}

impl Hash for HfmnCode {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.code.hash(state);
    }
}

impl Default for HfmnCode {
    fn default() -> Self {
        Self::new()
    }
}

impl Display for HfmnCode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let string = format!("{:064b}", self.code);
        let start_idx = string.len().saturating_sub(self.length as usize);
        write!(f, "{}", &string[start_idx..])
    }
}

// impl Index<usize> for Symbol {
//     type Output = Bit;
//     fn index(&self, index: usize) -> &Self::Output {
//         assert!(index < 64);

//         Bit::from_usize(self.symbol & (1 << index) >> index).unwrap()
//     }
// }

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

    #[test]
    fn direct_construction() {
        let input = 0b0101_1110;

        let symbol = HfmnCode::from(input, 8);

        assert_eq!(symbol.length, 8);
        assert_eq!(symbol.code, input);
    }

    #[test]
    fn push_construction() {
        let bits = vec![Bit::Zero, Bit::One, Bit::Zero, Bit::Zero];

        let mut symbol = HfmnCode::new();
        bits.into_iter().for_each(|b| symbol.push(b));

        assert_eq!(symbol.length, 4);
        assert_eq!(symbol.code, 0b0100);
    }

    #[test]
    fn retrieve_n_bits() {
        let bits = 0b11_0010;
        let symbol = HfmnCode::from(bits, 6);

        assert_eq!(symbol.first_n_bits(3), HfmnCode::from(bits & 0b111, 3));
        assert_eq!(symbol.first_n_bits(6), HfmnCode::from(bits, 6));
    }

    #[test]
    fn resize() {
        let mut symbol = HfmnCode::from(0b01_0101, 6);
        symbol.resize(8);
        assert_eq!(symbol.length, 8);
        assert_eq!(symbol.code, 0b0101_0100);
    }

    #[test]
    fn bits() {
        let symbol = HfmnCode::from(0b01_1001, 6);
        let expected = vec![
            Bit::Zero,
            Bit::One,
            Bit::One,
            Bit::Zero,
            Bit::Zero,
            Bit::One,
        ];
        assert_eq!(symbol.bits(), expected);
    }
}