chdlady-huffman 0.1.0

//! Canonical Huffman tree encoder with frequency-based code assignment.
use crate::bitstream::BitWriter;
use crate::error::HuffmanError;

#[derive(Clone, Copy, Debug, Default)]
struct BuildNode {
    parent: Option<usize>,
    weight: u32,
    bits: u32,
    numbits: u8,
}

/// Encoder for canonical Huffman bitstreams compatible with MAME CHD.
#[derive(Debug, Clone)]
pub struct HuffmanEncoder<const NUM_CODES: usize, const MAX_BITS: usize> {
    histogram: [u32; NUM_CODES],
    code_lengths: [u8; NUM_CODES],
    canonical_codes: [u32; NUM_CODES],
}

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

impl<const NUM_CODES: usize, const MAX_BITS: usize> HuffmanEncoder<NUM_CODES, MAX_BITS> {
    /// Creates a new encoder instance with an empty histogram.
    ///
    /// # Panics
    /// Panics if `MAX_BITS > 24`.
    pub fn new() -> Self {
        assert!(
            MAX_BITS <= 24,
            "MAX_BITS cannot exceed 24 in MAME Huffman format"
        );
        Self {
            histogram: [0; NUM_CODES],
            code_lengths: [0; NUM_CODES],
            canonical_codes: [0; NUM_CODES],
        }
    }

    /// Resets symbol counts in the histogram.
    pub fn histo_reset(&mut self) {
        self.histogram.fill(0);
    }

    /// Increments the frequency count of the given symbol.
    pub fn histo_one(&mut self, data: usize) {
        if data < NUM_CODES {
            self.histogram[data] = self.histogram[data].saturating_add(1);
        }
    }

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

    fn build_tree(&self, totaldata: u64, totalweight: u64) -> (usize, Vec<BuildNode>) {
        let mut nodes = vec![BuildNode::default(); NUM_CODES * 2];
        let mut active = Vec::with_capacity(NUM_CODES);

        for (curcode, node) in nodes.iter_mut().enumerate().take(NUM_CODES) {
            let count = self.histogram[curcode];
            if count != 0 {
                let mut w = ((count as u64) * totalweight / totaldata) as u32;
                if w == 0 {
                    w = 1;
                }
                *node = BuildNode {
                    parent: None,
                    weight: w,
                    bits: curcode as u32,
                    numbits: 0,
                };
                active.push(curcode);
            }
        }

        active.sort_by(|&a, &b| {
            nodes[b]
                .weight
                .cmp(&nodes[a].weight)
                .then_with(|| nodes[a].bits.cmp(&nodes[b].bits))
        });

        let mut nextalloc = NUM_CODES;
        let mut listitems = active.len();
        while listitems > 1 {
            let Some(node1_idx) = active.pop() else { break };
            let Some(node0_idx) = active.pop() else { break };
            listitems -= 2;

            let new_idx = nextalloc;
            nextalloc += 1;

            let combined_weight = nodes[node0_idx].weight + nodes[node1_idx].weight;
            nodes[new_idx] = BuildNode {
                parent: None,
                weight: combined_weight,
                bits: 0,
                numbits: 0,
            };
            nodes[node0_idx].parent = Some(new_idx);
            nodes[node1_idx].parent = Some(new_idx);

            let mut insert_pos = listitems;
            for i in 0..listitems {
                if combined_weight > nodes[active[i]].weight {
                    insert_pos = i;
                    break;
                }
            }
            active.insert(insert_pos, new_idx);
            listitems += 1;
        }

        let mut maxbits = 0usize;
        for curcode in 0..NUM_CODES {
            if nodes[curcode].weight > 0 {
                let mut bits = 0usize;
                let mut curr = Some(curcode);
                while let Some(idx) = curr {
                    if let Some(parent) = nodes[idx].parent {
                        bits += 1;
                        curr = Some(parent);
                    } else {
                        curr = None;
                    }
                }
                if bits == 0 {
                    bits = 1;
                }
                nodes[curcode].numbits = bits as u8;
                maxbits = maxbits.max(bits);
            }
        }

        (maxbits, nodes)
    }

    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(())
    }

    /// Computes optimal canonical Huffman tree bounded by MAX_BITS.
    pub fn compute_tree_from_histo(&mut self) -> Result<(), HuffmanError> {
        let sdatacount: u64 = self.histogram.iter().map(|&x| x as u64).sum();
        if sdatacount == 0 {
            self.code_lengths.fill(0);
            self.canonical_codes.fill(0);
            return Ok(());
        }

        let mut lowerweight = 0u64;
        let mut upperweight = sdatacount * 2;
        let mut best_nodes = vec![BuildNode::default(); NUM_CODES * 2];

        loop {
            let curweight = (upperweight + lowerweight) / 2;
            let (curmaxbits, candidate_nodes) = self.build_tree(sdatacount, curweight);

            if curmaxbits <= MAX_BITS {
                lowerweight = curweight;
                best_nodes = candidate_nodes;
                if curweight == sdatacount || (upperweight - lowerweight) <= 1 {
                    break;
                }
            } else {
                upperweight = curweight;
            }
        }

        for (len, best_node) in self.code_lengths.iter_mut().zip(&best_nodes) {
            *len = best_node.numbits;
        }

        self.assign_canonical_codes()
    }

    fn write_rle_tree_bits(bitbuf: &mut BitWriter, value: u8, mut repcount: usize, numbits: usize) {
        while repcount > 0 {
            if value == 1 {
                bitbuf.write(1, numbits);
                bitbuf.write(1, numbits);
                repcount -= 1;
            } else if repcount <= 2 {
                bitbuf.write(u32::from(value), numbits);
                repcount -= 1;
            } else {
                let cur_reps = (repcount - 3).min((1 << numbits) - 1);
                bitbuf.write(1, numbits);
                bitbuf.write(u32::from(value), numbits);
                bitbuf.write(cur_reps as u32, numbits);
                repcount -= cur_reps + 3;
            }
        }
    }

    /// Exports the tree using run-length encoding.
    pub fn export_tree_rle(&self, bitbuf: &mut BitWriter) -> Result<(), HuffmanError> {
        let numbits = if MAX_BITS >= 16 {
            5
        } else if MAX_BITS >= 8 {
            4
        } else {
            3
        };

        let mut lastval = 0xffu8;
        let mut repcount = 0usize;
        for curcode in 0..NUM_CODES {
            let newval = self.code_lengths[curcode];
            if newval == lastval {
                repcount += 1;
            } else {
                if repcount != 0 {
                    Self::write_rle_tree_bits(bitbuf, lastval, repcount, numbits);
                }
                lastval = newval;
                repcount = 1;
            }
        }

        if repcount != 0 {
            Self::write_rle_tree_bits(bitbuf, lastval, repcount, numbits);
        }

        Ok(())
    }

    /// Exports the tree using small Huffman compression.
    pub fn export_tree_huffman(&self, bitbuf: &mut BitWriter) -> Result<(), HuffmanError> {
        let mut rle_data = Vec::with_capacity(NUM_CODES);
        let mut rle_lengths = Vec::with_capacity(NUM_CODES / 3 + 1);
        let mut last = 0xffu8;
        let mut repcount = 0usize;

        let mut smallhuff = HuffmanEncoder::<24, 6>::new();

        for curcode in 0..NUM_CODES {
            let newval = self.code_lengths[curcode];
            if newval != last && repcount > 0 {
                if repcount == 1 {
                    let val = last + 1;
                    rle_data.push(val);
                    smallhuff.histo_one(val as usize);
                } else {
                    rle_data.push(0);
                    smallhuff.histo_one(0);
                    rle_lengths.push((repcount - 2) as u16);
                }
            }

            if newval == last {
                repcount += 1;
            } else {
                let val = newval + 1;
                rle_data.push(val);
                smallhuff.histo_one(val as usize);
                last = newval;
                repcount = 0;
            }
        }

        if repcount > 0 {
            if repcount == 1 {
                let val = last + 1;
                rle_data.push(val);
                smallhuff.histo_one(val as usize);
            } else {
                rle_data.push(0);
                smallhuff.histo_one(0);
                rle_lengths.push((repcount - 2) as u16);
            }
        }

        smallhuff.compute_tree_from_histo()?;

        let mut first_non_zero = 31usize;
        let mut last_non_zero = 0usize;
        for index in 1..24 {
            if smallhuff.code_lengths[index] != 0 {
                if first_non_zero == 31 {
                    first_non_zero = index;
                }
                last_non_zero = index;
            }
        }
        first_non_zero = first_non_zero.min(8);

        bitbuf.write(u32::from(smallhuff.code_lengths[0]), 3);
        bitbuf.write((first_non_zero.saturating_sub(1)) as u32, 3);
        for index in first_non_zero..=last_non_zero {
            bitbuf.write(u32::from(smallhuff.code_lengths[index]), 3);
        }
        bitbuf.write(7, 3);

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

        let mut length_idx = 0;
        for &data in &rle_data {
            smallhuff.encode_one(bitbuf, data as usize);
            if data == 0 {
                let count = rle_lengths[length_idx];
                length_idx += 1;
                if count < 7 {
                    bitbuf.write(u32::from(count), 3);
                } else {
                    bitbuf.write(7, 3);
                    bitbuf.write(u32::from(count - 7), rlefullbits);
                }
            }
        }

        Ok(())
    }

    /// Encodes a single symbol into the bitstream.
    pub fn encode_one(&self, bitbuf: &mut BitWriter, data: usize) {
        let code = self.canonical_codes[data];
        let numbits = self.code_lengths[data] as usize;
        bitbuf.write(code, numbits);
    }
}

/// Encodes an entire 8-bit buffer into a compressed byte vector using MAME canonical Huffman coding.
pub fn compress_huffman_8bit(source: &[u8]) -> Result<Vec<u8>, HuffmanError> {
    let mut encoder = HuffmanEncoder::<256, 16>::new();
    for &byte in source {
        encoder.histo_one(byte as usize);
    }
    encoder.compute_tree_from_histo()?;

    let mut writer = BitWriter::with_capacity(source.len());
    encoder.export_tree_huffman(&mut writer)?;
    for &byte in source {
        encoder.encode_one(&mut writer, byte as usize);
    }
    Ok(writer.into_bytes())
}