oxipkx 1.0.0

//! minimal raw-DEFLATE decompressor (RFC 1951), enough for ZIP method 8.
//!
//! ZIP stores DEFLATE streams in "raw" form: no zlib 2-byte header and no
//! trailing Adler checksum, just the bare compressed blocks. this is a
//! canonical-Huffman decoder in the style of Mark Adler's public-domain puff:
//! symbols are decoded by walking bit-lengths against per-length code counts.

/// error while inflating a DEFLATE stream.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InflateError {
    /// ran out of input bits before the stream terminated
    UnexpectedEnd,
    /// reserved block type 3, or other malformed block header
    BadBlockType,
    /// stored-block length did not match its one's-complement check
    BadStoredLength,
    /// a Huffman code did not resolve to any symbol
    BadCode,
    /// a back-reference pointed before the start of output
    BadDistance,
    /// produced more bytes than the entry's uncompressed size
    OutputOverflow,
}

const MAX_BITS: usize = 15;

/// length base values for symbols 257..285
const LENGTH_BASE: [u16; 29] = [
    3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131,
    163, 195, 227, 258,
];
/// extra bits appended to each length base
const LENGTH_EXTRA: [u8; 29] = [
    0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0,
];
/// distance base values for symbols 0..29
const DISTANCE_BASE: [u16; 30] = [
    1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537,
    2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577,
];
/// extra bits appended to each distance base
const DISTANCE_EXTRA: [u8; 30] = [
    0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13,
    13,
];
/// order in which code-length code lengths are stored in a dynamic block
const CODE_LENGTH_ORDER: [usize; 19] = [
    16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15,
];

/// canonical Huffman table: `counts[len]` codes of each bit length, and the
/// symbols sorted by (length, symbol) into `symbols`.
struct Huffman {
    counts: [u16; MAX_BITS + 1],
    symbols: Vec<u16>,
}

impl Huffman {
    /// build a table from a list of per-symbol code lengths (0 means unused).
    fn build(lengths: &[u8]) -> Self {
        let mut counts = [0u16; MAX_BITS + 1];
        for &length in lengths {
            counts[length as usize] += 1;
        }
        counts[0] = 0;
        let mut offsets = [0u16; MAX_BITS + 2];
        for length in 1..=MAX_BITS {
            offsets[length + 1] = offsets[length] + counts[length];
        }
        let mut symbols = vec![0u16; lengths.len()];
        for (symbol, &length) in lengths.iter().enumerate() {
            if length != 0 {
                symbols[offsets[length as usize] as usize] = symbol as u16;
                offsets[length as usize] += 1;
            }
        }
        Self { counts, symbols }
    }
}

/// LSB-first bit reader over the compressed byte stream.
struct BitReader<'a> {
    data: &'a [u8],
    byte_position: usize,
    bit_buffer: u32,
    bit_count: u32,
}

impl<'a> BitReader<'a> {
    fn new(data: &'a [u8]) -> Self {
        Self { data, byte_position: 0, bit_buffer: 0, bit_count: 0 }
    }

    /// read `count` bits (0..=32-ish), least-significant bit first.
    fn bits(&mut self, count: u32) -> Result<u32, InflateError> {
        while self.bit_count < count {
            let byte = *self.data.get(self.byte_position).ok_or(InflateError::UnexpectedEnd)?;
            self.byte_position += 1;
            self.bit_buffer |= (byte as u32) << self.bit_count;
            self.bit_count += 8;
        }
        let value = self.bit_buffer & ((1u32 << count) - 1);
        self.bit_buffer >>= count;
        self.bit_count -= count;
        Ok(value)
    }

    /// discard buffered bits back to a byte boundary (for stored blocks).
    fn align_to_byte(&mut self) {
        let drop = self.bit_count % 8;
        self.bit_buffer >>= drop;
        self.bit_count -= drop;
    }

    /// decode one symbol using a canonical Huffman table.
    fn decode(&mut self, table: &Huffman) -> Result<u16, InflateError> {
        let mut code = 0i32;
        let mut first = 0i32;
        let mut index = 0i32;
        for length in 1..=MAX_BITS {
            code |= self.bits(1)? as i32;
            let count = table.counts[length] as i32;
            if code - first < count {
                return Ok(table.symbols[(index + (code - first)) as usize]);
            }
            index += count;
            first = (first + count) << 1;
            code <<= 1;
        }
        Err(InflateError::BadCode)
    }
}

/// decompress a raw DEFLATE stream into exactly `expected_size` bytes.
pub fn inflate(input: &[u8], expected_size: usize) -> Result<Vec<u8>, InflateError> {
    let mut reader = BitReader::new(input);
    let mut output = Vec::with_capacity(expected_size);
    loop {
        let is_final = reader.bits(1)?;
        let block_type = reader.bits(2)?;
        match block_type {
            0 => inflate_stored(&mut reader, &mut output, expected_size)?,
            1 => inflate_block(&mut reader, &mut output, &fixed_tables(), expected_size)?,
            2 => {
                let (literals, distances) = read_dynamic_tables(&mut reader)?;
                inflate_block_with(&mut reader, &mut output, &literals, &distances, expected_size)?
            }
            _ => return Err(InflateError::BadBlockType),
        }
        if is_final == 1 {
            break;
        }
    }
    Ok(output)
}

fn inflate_stored(
    reader: &mut BitReader,
    output: &mut Vec<u8>,
    expected_size: usize,
) -> Result<(), InflateError> {
    reader.align_to_byte();
    let length = reader.bits(16)? as usize;
    let complement = reader.bits(16)? as usize;
    if length != (!complement & 0xffff) {
        return Err(InflateError::BadStoredLength);
    }
    for _ in 0..length {
        if output.len() >= expected_size {
            return Err(InflateError::OutputOverflow);
        }
        output.push(reader.bits(8)? as u8);
    }
    Ok(())
}

/// the fixed literal/length and distance tables (block type 1).
fn fixed_tables() -> (Huffman, Huffman) {
    let mut literal_lengths = [0u8; 288];
    for (symbol, length) in literal_lengths.iter_mut().enumerate() {
        *length = match symbol {
            0..=143 => 8,
            144..=255 => 9,
            256..=279 => 7,
            _ => 8,
        };
    }
    let distance_lengths = [5u8; 30];
    (Huffman::build(&literal_lengths), Huffman::build(&distance_lengths))
}

/// wrapper so the fixed-table path reads like the dynamic one.
fn inflate_block(
    reader: &mut BitReader,
    output: &mut Vec<u8>,
    tables: &(Huffman, Huffman),
    expected_size: usize,
) -> Result<(), InflateError> {
    inflate_block_with(reader, output, &tables.0, &tables.1, expected_size)
}

/// parse the dynamic block header into literal/length and distance tables.
fn read_dynamic_tables(reader: &mut BitReader) -> Result<(Huffman, Huffman), InflateError> {
    let literal_count = reader.bits(5)? as usize + 257;
    let distance_count = reader.bits(5)? as usize + 1;
    let code_length_count = reader.bits(4)? as usize + 4;

    let mut code_length_lengths = [0u8; 19];
    for index in 0..code_length_count {
        code_length_lengths[CODE_LENGTH_ORDER[index]] = reader.bits(3)? as u8;
    }
    let code_length_table = Huffman::build(&code_length_lengths);

    let total = literal_count + distance_count;
    let mut lengths = vec![0u8; total];
    let mut index = 0;
    while index < total {
        let symbol = reader.decode(&code_length_table)?;
        match symbol {
            0..=15 => {
                lengths[index] = symbol as u8;
                index += 1;
            }
            16 => {
                // repeat previous length 3..6 times
                if index == 0 {
                    return Err(InflateError::BadCode);
                }
                let previous = lengths[index - 1];
                let repeat = reader.bits(2)? as usize + 3;
                for _ in 0..repeat {
                    if index >= total {
                        return Err(InflateError::BadCode);
                    }
                    lengths[index] = previous;
                    index += 1;
                }
            }
            17 => {
                // repeat zero 3..10 times
                let repeat = reader.bits(3)? as usize + 3;
                index = fill_zeros(&mut lengths, index, repeat)?;
            }
            18 => {
                // repeat zero 11..138 times
                let repeat = reader.bits(7)? as usize + 11;
                index = fill_zeros(&mut lengths, index, repeat)?;
            }
            _ => return Err(InflateError::BadCode),
        }
    }

    let literals = Huffman::build(&lengths[..literal_count]);
    let distances = Huffman::build(&lengths[literal_count..]);
    Ok((literals, distances))
}

fn fill_zeros(lengths: &mut [u8], mut index: usize, repeat: usize) -> Result<usize, InflateError> {
    for _ in 0..repeat {
        if index >= lengths.len() {
            return Err(InflateError::BadCode);
        }
        lengths[index] = 0;
        index += 1;
    }
    Ok(index)
}

/// decode compressed data using the given literal/length and distance tables.
fn inflate_block_with(
    reader: &mut BitReader,
    output: &mut Vec<u8>,
    literals: &Huffman,
    distances: &Huffman,
    expected_size: usize,
) -> Result<(), InflateError> {
    loop {
        let symbol = reader.decode(literals)?;
        match symbol {
            0..=255 => {
                if output.len() >= expected_size {
                    return Err(InflateError::OutputOverflow);
                }
                output.push(symbol as u8);
            }
            256 => return Ok(()), // end of block
            257..=285 => {
                let length_symbol = (symbol - 257) as usize;
                let length = LENGTH_BASE[length_symbol] as usize
                    + reader.bits(LENGTH_EXTRA[length_symbol] as u32)? as usize;
                let distance_symbol = reader.decode(distances)? as usize;
                if distance_symbol >= DISTANCE_BASE.len() {
                    return Err(InflateError::BadCode);
                }
                let distance = DISTANCE_BASE[distance_symbol] as usize
                    + reader.bits(DISTANCE_EXTRA[distance_symbol] as u32)? as usize;
                if distance > output.len() {
                    return Err(InflateError::BadDistance);
                }
                let start = output.len() - distance;
                for offset in 0..length {
                    if output.len() >= expected_size {
                        return Err(InflateError::OutputOverflow);
                    }
                    output.push(output[start + offset]);
                }
            }
            _ => return Err(InflateError::BadCode),
        }
    }
}