zstd_zero/

huffman.rs

use crate::bitstream::BackwardBits;
use crate::fse::Table as FseTable;
use crate::DecodeError;

const MAX_BITS: u8 = 11;

/// An initialized entropy-table slot; the decoder manages its contents.
#[derive(Clone, Copy, Debug, Default)]
pub struct Entry {
    symbol: u8,
    bits: u8,
}

impl Entry {
    /// Create an empty slot, including for statically allocated workspaces.
    pub const fn new() -> Self {
        Self { symbol: 0, bits: 0 }
    }
}

pub(crate) struct Table<'a> {
    entries: &'a mut [Entry],
    scratch: FseTable<'a>,
    table_bits: u8,
    valid: bool,
}

impl<'a> Table<'a> {
    pub(crate) fn new(entries: &'a mut [Entry], scratch: FseTable<'a>) -> Self {
        Self {
            entries,
            scratch,
            table_bits: 0,
            valid: false,
        }
    }

    pub(crate) fn reset(&mut self) {
        self.table_bits = 0;
        self.valid = false;
        self.scratch.reset();
    }

    pub(crate) fn is_valid(&self) -> bool {
        self.valid
    }

    pub(crate) fn read_description(
        &mut self,
        input: &[u8],
        fse_scratch: &mut [i16],
    ) -> Result<usize, DecodeError> {
        let header = *input.first().ok_or(DecodeError::InvalidEntropyTable)?;
        let mut weights = [0u8; 256];
        let (weight_count, consumed) = if header < 128 {
            let compressed_size = header as usize;
            if compressed_size == 0 || input.len() < 1 + compressed_size {
                return Err(DecodeError::InvalidEntropyTable);
            }
            let count = decode_compressed_weights(
                &input[1..1 + compressed_size],
                &mut weights,
                &mut self.scratch,
                fse_scratch,
            )?;
            (count, 1 + compressed_size)
        } else {
            let count = header as usize - 127;
            let bytes = count.div_ceil(2);
            if input.len() < 1 + bytes {
                return Err(DecodeError::InvalidEntropyTable);
            }
            for (index, weight) in weights[..count].iter_mut().enumerate() {
                let packed = input[1 + index / 2];
                *weight = if index % 2 == 0 {
                    packed >> 4
                } else {
                    packed & 0x0f
                };
            }
            (count, 1 + bytes)
        };
        self.build(&weights[..weight_count])?;
        Ok(consumed)
    }

    pub(crate) fn decode(
        &self,
        input: &[u8],
        output: &mut [u8],
        strict: bool,
    ) -> Result<(), DecodeError> {
        if !self.valid {
            return Err(DecodeError::InvalidEntropyTable);
        }
        let mut bits = BackwardBits::new(input)?;
        for byte in output {
            let index = bits.peek_padded(self.table_bits) as usize;
            let entry = self.entries[index];
            if entry.bits == 0 {
                return Err(DecodeError::InvalidEntropyTable);
            }
            let consumed = if strict {
                // A short read must fail here. Clamping to the bits that remain would
                // drive `remaining` to exactly 0, making the end-of-stream check below
                // vacuous, and every remaining literal would then be decoded from an
                // all-zero peek — silently emitting `entries[0]` instead of erroring.
                entry.bits
            } else {
                core::cmp::min(entry.bits as usize, bits.remaining()) as u8
            };
            bits.consume(consumed)?;
            *byte = entry.symbol;
        }
        if bits.remaining() != 0 {
            return Err(DecodeError::InvalidBitstream);
        }
        Ok(())
    }

    pub(crate) fn decode_four(
        &self,
        input: &[u8],
        output: &mut [u8],
        strict: bool,
    ) -> Result<(), DecodeError> {
        if input.len() < 6 {
            return Err(DecodeError::InvalidBitstream);
        }
        let sizes = [
            u16::from_le_bytes([input[0], input[1]]) as usize,
            u16::from_le_bytes([input[2], input[3]]) as usize,
            u16::from_le_bytes([input[4], input[5]]) as usize,
        ];
        let mut starts = [6usize; 4];
        for (stream, size) in sizes.iter().enumerate() {
            starts[stream + 1] = starts[stream]
                .checked_add(*size)
                .ok_or(DecodeError::ArithmeticOverflow)?;
        }
        if starts[3] > input.len() {
            return Err(DecodeError::InvalidBitstream);
        }
        let ends = [starts[1], starts[2], starts[3], input.len()];
        let segment = output.len().div_ceil(4);
        for stream in 0..4 {
            let output_start = core::cmp::min(stream * segment, output.len());
            let output_end = core::cmp::min(output_start + segment, output.len());
            if output_start != output_end {
                self.decode(
                    &input[starts[stream]..ends[stream]],
                    &mut output[output_start..output_end],
                    strict,
                )?;
            }
        }
        Ok(())
    }

    fn build(&mut self, weights: &[u8]) -> Result<(), DecodeError> {
        if weights.is_empty() || weights.len() >= 256 {
            return Err(DecodeError::InvalidEntropyTable);
        }
        let mut sum = 0u32;
        for weight in weights {
            if *weight > MAX_BITS {
                return Err(DecodeError::InvalidEntropyTable);
            }
            if *weight != 0 {
                sum = sum
                    .checked_add(1u32 << (*weight - 1))
                    .ok_or(DecodeError::InvalidEntropyTable)?;
            }
        }
        if sum == 0 {
            return Err(DecodeError::InvalidEntropyTable);
        }
        let table_bits = 32 - sum.leading_zeros();
        if table_bits > MAX_BITS as u32 {
            return Err(DecodeError::InvalidEntropyTable);
        }
        let remainder = (1u32 << table_bits)
            .checked_sub(sum)
            .ok_or(DecodeError::InvalidEntropyTable)?;
        if !remainder.is_power_of_two() {
            return Err(DecodeError::InvalidEntropyTable);
        }
        let implied_weight = 32 - remainder.leading_zeros();
        let table_bits = table_bits as u8;

        let mut code_bits = [0u8; 256];
        for (index, weight) in weights.iter().enumerate() {
            code_bits[index] = if *weight == 0 {
                0
            } else {
                table_bits + 1 - *weight
            };
        }
        code_bits[weights.len()] = table_bits + 1 - implied_weight as u8;

        let mut ranks = [0usize; MAX_BITS as usize + 1];
        for bits in &code_bits[..=weights.len()] {
            ranks[*bits as usize] += 1;
        }
        let mut starts = [0usize; MAX_BITS as usize + 1];
        for bits in (1..=table_bits as usize).rev() {
            starts[bits - 1] =
                starts[bits] + ranks[bits] * (1usize << (table_bits as usize - bits));
        }
        let used = 1usize << table_bits;
        if starts[0] != used {
            return Err(DecodeError::InvalidEntropyTable);
        }
        self.entries[..used].fill(Entry::default());
        for (symbol, bits) in code_bits[..=weights.len()].iter().enumerate() {
            if *bits == 0 {
                continue;
            }
            let width = 1usize << (table_bits - *bits);
            let start = starts[*bits as usize];
            let end = start
                .checked_add(width)
                .ok_or(DecodeError::InvalidEntropyTable)?;
            if end > used {
                return Err(DecodeError::InvalidEntropyTable);
            }
            self.entries[start..end].fill(Entry {
                symbol: symbol as u8,
                bits: *bits,
            });
            starts[*bits as usize] = end;
        }
        self.table_bits = table_bits;
        self.valid = true;
        Ok(())
    }
}

fn decode_compressed_weights(
    input: &[u8],
    output: &mut [u8; 256],
    table: &mut FseTable<'_>,
    scratch: &mut [i16],
) -> Result<usize, DecodeError> {
    table.reset();
    let description_size = table.read_description(input, 12, 6, scratch)?;
    if description_size >= input.len() {
        return Err(DecodeError::InvalidEntropyTable);
    }
    let mut bits = BackwardBits::new(&input[description_size..])?;
    let mut first = bits.read(table.log())?;
    let mut second = bits.read(table.log())?;
    let mut count = 0usize;
    loop {
        if count >= 255 {
            return Err(DecodeError::InvalidEntropyTable);
        }
        output[count] = table.symbol(first)?;
        count += 1;
        let entry = table.entry(first)?;
        if bits.remaining() < entry.bits as usize {
            if count >= 255 {
                return Err(DecodeError::InvalidEntropyTable);
            }
            output[count] = table.symbol(second)?;
            return Ok(count + 1);
        }
        table.update(&mut first, &mut bits)?;
        core::mem::swap(&mut first, &mut second);
    }
}

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

    #[test]
    fn truncated_literal_symbol_is_rejected() {
        let mut entries = std::vec![Entry::new(); crate::HUFFMAN_ENTRIES];
        let mut scratch = std::vec![crate::FseEntry::new(); 512];
        let mut table = Table::new(&mut entries, FseTable::new(&mut scratch));
        table.build(&[1]).unwrap();
        let mut output = [0; 1];
        assert!(table.decode(&[1], &mut output, true).is_err());
        assert!(table.decode(&[2], &mut output, true).is_ok());
        assert_eq!(output, [0]);
    }
}