chdlady-huffman 0.1.0

//! Canonical Huffman tree decoder with fast lookup table.
use crate::bitstream::BitReader;
use crate::error::HuffmanError;

/// Decoder for canonical Huffman bitstreams compatible with MAME CHD.
#[derive(Debug, Clone)]
pub struct HuffmanDecoder<const NUM_CODES: usize, const MAX_BITS: usize> {
    pub(crate) code_lengths: [u8; NUM_CODES],
    pub(crate) canonical_codes: [u32; NUM_CODES],
    lookup: Vec<u16>,
}

impl<const NUM_CODES: usize, const MAX_BITS: usize> Default
    for HuffmanDecoder<NUM_CODES, MAX_BITS>
{
    fn default() -> Self {
        Self::new()
    }
}

impl<const NUM_CODES: usize, const MAX_BITS: usize> HuffmanDecoder<NUM_CODES, MAX_BITS> {
    /// Creates a new decoder instance with an empty code tree.
    ///
    /// # Panics
    /// Panics if `MAX_BITS > 24` or `NUM_CODES > 2048`.
    pub fn new() -> Self {
        assert!(
            MAX_BITS <= 24,
            "MAX_BITS cannot exceed 24 in MAME Huffman format"
        );
        assert!(
            NUM_CODES <= 2048,
            "NUM_CODES cannot exceed 2048 when using 16-bit lookup table entries"
        );
        Self {
            code_lengths: [0; NUM_CODES],
            canonical_codes: [0; NUM_CODES],
            lookup: vec![0; 1 << MAX_BITS],
        }
    }

    /// Returns the canonical code lengths.
    pub fn code_lengths(&self) -> &[u8; NUM_CODES] {
        &self.code_lengths
    }

    fn assign_canonical_codes(&mut self) -> Result<(), HuffmanError> {
        let mut bithisto = [0u32; 33];
        for &len in &self.code_lengths {
            let len = len as usize;
            if len > MAX_BITS {
                return Err(HuffmanError::InternalInconsistency);
            }
            if len <= 32 {
                bithisto[len] += 1;
            }
        }

        let mut curstart = 0u32;
        for codelen in (1..=32).rev() {
            let nextstart = (curstart + bithisto[codelen]) >> 1;
            if codelen != 1 && nextstart * 2 != (curstart + bithisto[codelen]) {
                return Err(HuffmanError::InternalInconsistency);
            }
            bithisto[codelen] = curstart;
            curstart = nextstart;
        }

        for i in 0..NUM_CODES {
            let len = self.code_lengths[i] as usize;
            if len > 0 {
                self.canonical_codes[i] = bithisto[len];
                bithisto[len] += 1;
            } else {
                self.canonical_codes[i] = 0;
            }
        }
        Ok(())
    }

    fn build_lookup_table(&mut self) {
        self.lookup.fill(0);
        for curcode in 0..NUM_CODES {
            let numbits = self.code_lengths[curcode] as usize;
            if numbits > 0 {
                let value = ((curcode as u16) << 5) | ((numbits as u16) & 0x1f);
                let shift = MAX_BITS - numbits;
                let start = (self.canonical_codes[curcode] as usize) << shift;
                let end = (((self.canonical_codes[curcode] + 1) as usize) << shift) - 1;
                for entry in &mut self.lookup[start..=end] {
                    *entry = value;
                }
            }
        }
    }

    /// Decodes a single symbol from the bitstream.
    #[inline(always)]
    pub fn decode_one(&self, bitbuf: &mut BitReader<'_>) -> u32 {
        let bits = bitbuf.peek(MAX_BITS) as usize;
        let lookup = self.lookup[bits];
        bitbuf.remove((lookup & 0x1f) as usize);
        u32::from(lookup >> 5)
    }

    /// Imports tree lengths encoded using run-length encoding.
    pub fn import_tree_rle(&mut self, bitbuf: &mut BitReader<'_>) -> Result<(), HuffmanError> {
        let numbits = if MAX_BITS >= 16 {
            5
        } else if MAX_BITS >= 8 {
            4
        } else {
            3
        };

        let mut curnode = 0;
        while curnode < NUM_CODES {
            let nodebits = bitbuf.read(numbits) as u8;
            if nodebits != 1 {
                self.code_lengths[curnode] = nodebits;
                curnode += 1;
            } else {
                let nodebits = bitbuf.read(numbits) as u8;
                if nodebits == 1 {
                    self.code_lengths[curnode] = 1;
                    curnode += 1;
                } else {
                    let repcount = (bitbuf.read(numbits) + 3) as usize;
                    for _ in 0..repcount {
                        if curnode >= NUM_CODES {
                            return Err(HuffmanError::InvalidData);
                        }
                        self.code_lengths[curnode] = nodebits;
                        curnode += 1;
                    }
                }
            }
        }

        if curnode != NUM_CODES {
            return Err(HuffmanError::InvalidData);
        }

        self.assign_canonical_codes()?;
        self.build_lookup_table();

        if bitbuf.overflow() {
            Err(HuffmanError::InputBufferTooSmall)
        } else {
            Ok(())
        }
    }

    /// Imports tree lengths encoded using a small Huffman sub-tree.
    pub fn import_tree_huffman(&mut self, bitbuf: &mut BitReader<'_>) -> Result<(), HuffmanError> {
        let mut smallhuff = HuffmanDecoder::<24, 6>::new();
        smallhuff.code_lengths[0] = bitbuf.read(3) as u8;
        let start = (bitbuf.read(3) + 1) as usize;
        let mut count = 0;
        for index in 1..24 {
            if index < start || count == 7 {
                smallhuff.code_lengths[index] = 0;
            } else {
                count = bitbuf.read(3);
                smallhuff.code_lengths[index] = if count == 7 { 0 } else { count as u8 };
            }
        }

        smallhuff.assign_canonical_codes()?;
        smallhuff.build_lookup_table();

        let mut temp = (NUM_CODES.saturating_sub(9)) as u32;
        let mut rlefullbits = 0usize;
        while temp != 0 {
            temp >>= 1;
            rlefullbits += 1;
        }

        let mut last = 0u8;
        let mut curcode = 0;
        while curcode < NUM_CODES {
            let value = smallhuff.decode_one(bitbuf);
            if value != 0 {
                last = (value - 1) as u8;
                self.code_lengths[curcode] = last;
                curcode += 1;
            } else {
                let mut count = (bitbuf.read(3) + 2) as usize;
                if count == 7 + 2 {
                    count += bitbuf.read(rlefullbits) as usize;
                }
                while count != 0 && curcode < NUM_CODES {
                    self.code_lengths[curcode] = last;
                    curcode += 1;
                    count -= 1;
                }
            }
        }

        if curcode != NUM_CODES {
            return Err(HuffmanError::InvalidData);
        }

        self.assign_canonical_codes()?;
        self.build_lookup_table();

        if bitbuf.overflow() {
            Err(HuffmanError::InputBufferTooSmall)
        } else {
            Ok(())
        }
    }

    /// Sets explicit code lengths and rebuilds the lookup table.
    pub fn set_code_lengths(&mut self, lengths: [u8; NUM_CODES]) -> Result<(), HuffmanError> {
        self.code_lengths = lengths;
        self.assign_canonical_codes()?;
        self.build_lookup_table();
        Ok(())
    }
}