#![warn(clippy::all)]

//! EFI/Tiano firmware decompression.
//!
//! This is a clean Rust port of EDK2's `BaseTools/Source/C/Common/Decompress.c`.
//! Supports both the EFI Standard and Tiano compression variants.
//!
//! # Header format
//!
//! The 8-byte header is:
//! - bytes 0..4: CompressedSize (u32 LE) — size of compressed payload after the header
//! - bytes 4..8: OriginalSize   (u32 LE) — size of decompressed output
//!
//! # Algorithm overview
//!
//! The algorithm is an LZSS + adaptive Huffman scheme. Three canonical Huffman
//! tables are maintained: the C-table (char/length codes), the PT-table (position
//! and extra-set codes). Blocks are decoded one at a time; each block begins with
//! fresh table definitions.

use std::fmt;

// ---------------------------------------------------------------------------
// Constants (verbatim from Decompress.c)
// ---------------------------------------------------------------------------

/// Bit-buffer size in bits (u32 holds the buffer).
const BITBUFSIZ: u16 = 32;

/// Maximum match length for back-references.
const MAXMATCH: u32 = 256;

/// Minimum match length for back-references.
const THRESHOLD: u32 = 3;

/// Number of bits in a code unit (used to derive NT).
const CODE_BIT: u16 = 16;

/// Number of symbols in the char/length set.
/// NC = 0xFF + MAXMATCH + 2 - THRESHOLD = 255 + 256 + 2 - 3 = 510
const NC: usize = (0xFF_u32 + MAXMATCH + 2 - THRESHOLD) as usize; // 510

/// Bits needed to index the C decode table (table size = 4096 = 2^12).
const CBIT: u16 = 9;

/// Number of position-length codes for the extra set.
/// NT = CODE_BIT + 3 = 19
const NT: usize = (CODE_BIT + 3) as usize; // 19

/// Bits used to encode the NT count.
const TBIT: u16 = 5;

/// Maximum position bits for Tiano.
const MAXPBIT: u16 = 5;

/// Maximum position bits for EFI Standard.
const EFIPBIT: u16 = 4;

/// Maximum value of a position code (MAXNP = (1 << MAXPBIT) - 1 = 31).
const MAXNP: usize = (1_usize << (MAXPBIT as usize)) - 1; // 31

/// NPT = max(NT, MAXNP) = max(19, 31) = 31
const NPT: usize = if NT > MAXNP { NT } else { MAXNP }; // 31

/// Sentinel returned by MakeTable on corruption.
const BAD_TABLE: u16 = u16::MAX;

/// Maximum decompressed size accepted from a header.
///
/// Hostiles can claim an arbitrarily large `orig_size` to force a multi-GB
/// allocation.  64 MiB is a generous upper bound for any real UEFI firmware
/// section.
const MAX_DECOMPRESS_SIZE: u32 = 64 * 1024 * 1024;

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Decoded EFI compression header.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EfiHeader {
    /// Size of the compressed payload, excluding the 8-byte header.
    pub compressed_size: u32,
    /// Expected size of the decompressed output.
    pub original_size: u32,
}

/// Which compression variant to use for decompression.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CompressionType {
    /// EFI Standard compression (position bits = 4).
    EfiStandard,
    /// Tiano compression (position bits = 5).
    Tiano,
}

/// Errors that can occur during decompression.
#[derive(Debug, PartialEq, Eq)]
pub enum DecompressError {
    /// The input slice is shorter than required.
    InputTooShort,
    /// The compressed data is malformed.
    InvalidData(&'static str),
}

impl fmt::Display for DecompressError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DecompressError::InputTooShort => write!(f, "input too short"),
            DecompressError::InvalidData(msg) => write!(f, "invalid data: {msg}"),
        }
    }
}

impl std::error::Error for DecompressError {}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Parse the 8-byte header and return an [`EfiHeader`].
///
/// `src` must be at least 8 bytes; the header layout is:
/// - `[0..4]`: CompressedSize (u32 LE) — size of the compressed payload after the header
/// - `[4..8]`: OriginalSize   (u32 LE) — size of the decompressed output
pub fn parse_header(src: &[u8]) -> Result<EfiHeader, DecompressError> {
    if src.len() < 8 {
        return Err(DecompressError::InputTooShort);
    }
    let compressed_size = u32::from_le_bytes(src[0..4].try_into().unwrap());
    let original_size = u32::from_le_bytes(src[4..8].try_into().unwrap());

    // Validate: total source must be at least header + compressed payload.
    // Also guard against wrap: compressed_size + 8 must not overflow and must fit in src.
    let total = (compressed_size as usize).saturating_add(8);
    if src.len() < total {
        return Err(DecompressError::InvalidData("compressed size exceeds source buffer"));
    }

    Ok(EfiHeader { compressed_size, original_size })
}

/// Decompress EFI/Tiano compressed data.
///
/// `src` must include the 8-byte header followed by the compressed payload.
/// Returns the decompressed bytes on success.
pub fn decompress(src: &[u8], compression_type: CompressionType) -> Result<Vec<u8>, DecompressError> {
    let header = parse_header(src)?;
    let orig_size = header.original_size;
    let comp_size = header.compressed_size;

    if orig_size == 0 {
        return Ok(Vec::new());
    }

    if orig_size > MAX_DECOMPRESS_SIZE {
        return Err(DecompressError::InvalidData("original size exceeds maximum (64 MiB)"));
    }

    let pbit = match compression_type {
        CompressionType::EfiStandard => EFIPBIT,
        CompressionType::Tiano => MAXPBIT,
    };

    let payload = &src[8..8 + comp_size as usize];
    let decoder = Decoder::new(payload, orig_size, pbit);
    decoder.decode()
}

// ---------------------------------------------------------------------------
// Decoder state
// ---------------------------------------------------------------------------

struct Decoder<'a> {
    /// Compressed source bytes (payload only, after header).
    src: &'a [u8],
    /// Current read position into `src`.
    in_buf: usize,

    /// 32-bit bit buffer: top bits are the next bits to consume.
    bit_buf: u32,
    /// Staging byte for partial reads.
    sub_bit_buf: u32,
    /// How many valid bits remain in `sub_bit_buf`.
    bit_count: u16,

    /// Remaining bytes of compressed source.
    comp_size: u32,

    /// Expected original (decompressed) size.
    orig_size: u32,

    /// Output buffer (also serves as the sliding window).
    dst: Vec<u8>,
    /// Write position in `dst`.
    out_buf: u32,

    /// Remaining symbols in the current block.
    block_size: u16,

    /// Set non-zero if a table decode error occurs.
    bad_table_flag: u16,

    /// Position-bit width (4 for EFI Standard, 5 for Tiano).
    pbit: u16,

    // Huffman decode tables
    /// C decode table (4096 entries, indexed by top 12 bits of bit_buf).
    c_table: Box<[u16; 4096]>,
    /// PT decode table (256 entries, indexed by top 8 bits of bit_buf).
    pt_table: Box<[u16; 256]>,

    // Code-length arrays
    c_len: Box<[u8; NC]>,
    pt_len: Box<[u8; NPT]>,

    // Binary tree for long codes (indices relative to NC for C, MAXNP for PT).
    left: Box<[u16; 2 * NC - 1]>,
    right: Box<[u16; 2 * NC - 1]>,
}

impl<'a> Decoder<'a> {
    fn new(src: &'a [u8], orig_size: u32, pbit: u16) -> Self {
        Decoder {
            src,
            in_buf: 0,
            bit_buf: 0,
            sub_bit_buf: 0,
            bit_count: 0,
            comp_size: src.len() as u32,
            orig_size,
            dst: vec![0u8; orig_size as usize],
            out_buf: 0,
            block_size: 0,
            bad_table_flag: 0,
            pbit,
            c_table: Box::new([0u16; 4096]),
            pt_table: Box::new([0u16; 256]),
            c_len: Box::new([0u8; NC]),
            pt_len: Box::new([0u8; NPT]),
            left: Box::new([0u16; 2 * NC - 1]),
            right: Box::new([0u16; 2 * NC - 1]),
        }
    }

    // -----------------------------------------------------------------------
    // Bit I/O
    // -----------------------------------------------------------------------

    /// Shift `bit_buf` left by `n` bits, refilling from `sub_bit_buf` and source.
    ///
    /// Mirrors `FillBuf()` in Decompress.c exactly.
    fn fill_buf(&mut self, mut n: u16) {
        // Shift the bit buffer left by n, using u64 arithmetic to avoid overflow.
        self.bit_buf = (((self.bit_buf as u64) << n) & 0xFFFF_FFFF) as u32;

        while n > self.bit_count {
            // OR in the sub_bit_buf shifted appropriately.
            n -= self.bit_count;
            self.bit_buf |= (((self.sub_bit_buf as u64) << n) & 0xFFFF_FFFF) as u32;

            if self.comp_size > 0 {
                self.comp_size -= 1;
                self.sub_bit_buf = self.src[self.in_buf] as u32;
                self.in_buf += 1;
                self.bit_count = 8;
            } else {
                self.sub_bit_buf = 0;
                self.bit_count = 8;
            }
        }

        self.bit_count -= n;
        self.bit_buf |= self.sub_bit_buf >> self.bit_count;
    }

    /// Return the top `n` bits of `bit_buf`, then call `fill_buf(n)`.
    ///
    /// Mirrors `GetBits()` in Decompress.c.
    fn get_bits(&mut self, n: u16) -> u32 {
        let out = self.bit_buf >> (BITBUFSIZ - n);
        self.fill_buf(n);
        out
    }

    // -----------------------------------------------------------------------
    // Table construction
    // -----------------------------------------------------------------------

    /// Build a canonical Huffman decode table from code lengths.
    ///
    /// Mirrors `MakeTable()` in Decompress.c.
    ///
    /// # Parameters
    /// - `num_of_char`: number of symbols
    /// - `bit_len`: slice of code lengths, length `num_of_char`
    /// - `table_bits`: log2 of the direct lookup table size
    /// - `table`: the output table
    ///
    /// Returns `BAD_TABLE` on corruption, otherwise 0.
    fn make_table(
        left: &mut [u16; 2 * NC - 1],
        right: &mut [u16; 2 * NC - 1],
        num_of_char: u16,
        bit_len: &[u8],
        table_bits: u16,
        table: &mut [u16],
    ) -> u16 {
        let mut count = [0u16; 17];
        let mut weight = [0u16; 17];
        let mut start = [0u16; 18];

        // count[1..=16] start zeroed; these loops use the index as arithmetic input,
        // not merely as an accessor, so needless_range_loop does not apply here.
        #[allow(clippy::needless_range_loop)]
        for idx in 1..=16usize {
            count[idx] = 0;
        }
        for (idx, &bl) in bit_len[..num_of_char as usize].iter().enumerate() {
            if bl > 16 {
                return BAD_TABLE;
            }
            let _ = idx; // `idx` not used here; pattern satisfies clippy
            count[bl as usize] += 1;
        }

        start[1] = 0;
        #[allow(clippy::needless_range_loop)]
        for idx in 1..=16usize {
            start[idx + 1] = start[idx].wrapping_add(count[idx] << (16 - idx));
        }

        // start[17] must be 0 (i.e. == 1<<16 mod 2^16, which wraps to 0).
        if start[17] != 0 {
            return BAD_TABLE;
        }

        let ju_bits = 16 - table_bits;

        // These loops use `idx` as an arithmetic operand (shift amount), not only as
        // an array index, so the needless_range_loop suggestion would be incorrect.
        #[allow(clippy::needless_range_loop)]
        for idx in 1..=(table_bits as usize) {
            start[idx] >>= ju_bits;
            weight[idx] = 1u16 << (table_bits - idx as u16);
        }
        #[allow(clippy::needless_range_loop)]
        for idx in (table_bits as usize + 1)..=16 {
            weight[idx] = 1u16 << (16 - idx as u16);
        }

        // Zero-fill the tail of the direct table not covered by any code.
        let mut idx = (start[table_bits as usize + 1] >> ju_bits) as usize;
        if idx != 0 {
            let table_len = 1usize << table_bits;
            while idx < table_len {
                table[idx] = 0;
                idx += 1;
            }
        }

        let mut avail = num_of_char as usize;
        let mask: u16 = 1u16 << (15 - table_bits);
        let max_table_length: u16 = 1u16 << table_bits;

        // The C source uses a raw pointer that walks between `table`, `left`, and
        // `right` depending on bit values. We model the current "pointer location"
        // as a tagged location: either a slot in `table` or a slot in `left`/`right`.
        // InLeft/InRight carry the node index into their respective arrays.
        enum Loc {
            Table(usize),
            Left(usize),
            Right(usize),
        }

        // Read the current cell, whichever array it lives in.
        let read_loc = |loc: &Loc, table: &[u16], left: &[u16], right: &[u16]| -> u16 {
            match loc {
                Loc::Table(i) => table[*i],
                Loc::Left(i) => left[*i],
                Loc::Right(i) => right[*i],
            }
        };

        // Write a value to the current cell.
        let write_loc =
            |loc: &Loc, val: u16, table: &mut [u16], left: &mut [u16], right: &mut [u16]| {
                match loc {
                    Loc::Table(i) => table[*i] = val,
                    Loc::Left(i) => left[*i] = val,
                    Loc::Right(i) => right[*i] = val,
                }
            };

        // `ch` is written as the symbol value into the table, so we need it as a
        // numeric index, not just to index bit_len. The enumerate rewrite clippy
        // suggests would work but `ch` is the symbol, not merely a counter.
        #[allow(clippy::needless_range_loop)]
        for ch in 0..num_of_char as usize {
            let len = bit_len[ch];
            if len == 0 || len >= 17 {
                continue;
            }
            let len = len as usize;

            let next_code = start[len].wrapping_add(weight[len]);

            if len <= table_bits as usize {
                // Direct table fill.
                if start[len] >= next_code || next_code > max_table_length {
                    return BAD_TABLE;
                }
                for i in start[len]..next_code {
                    table[i as usize] = ch as u16;
                }
            } else {
                // Long code: walk/build the binary tree in left[]/right[].
                let mut idx3 = start[len];
                // The C `Pointer` starts pointing at `Table[idx3 >> JuBits]`.
                let mut loc = Loc::Table((idx3 >> ju_bits) as usize);
                let mut depth = len - table_bits as usize;
                let mut bit_mask = mask;

                while depth != 0 {
                    let cur = read_loc(&loc, table, left, right);
                    let node = if cur == 0 {
                        // Allocate a new tree node.
                        left[avail] = 0;
                        right[avail] = 0;
                        write_loc(&loc, avail as u16, table, left, right);
                        let n = avail;
                        avail += 1;
                        n
                    } else {
                        cur as usize
                    };

                    // Follow the right or left branch depending on the current bit.
                    loc = if idx3 & bit_mask != 0 {
                        Loc::Right(node)
                    } else {
                        Loc::Left(node)
                    };

                    idx3 <<= 1;
                    depth -= 1;
                    bit_mask >>= 1;
                }

                write_loc(&loc, ch as u16, table, left, right);
            }

            start[len] = next_code;
        }

        0
    }

    /// Decode a single canonical-Huffman symbol.
    ///
    /// Index the direct `table` by the top `index_bits` of `bit_buf`; if the
    /// result is a tree pointer (`>= threshold`), walk `left`/`right` by
    /// successive bits until a leaf symbol (`< threshold`) is reached.
    ///
    /// Shared by `read_c_len` (PT table), `decode_c` (C table), and `decode_p`
    /// (PT table) — the three call sites differ only in table, index width, and
    /// threshold. Does not consume bits; callers `fill_buf` the symbol's length.
    fn decode_symbol(&self, table: &[u16], index_bits: u16, threshold: u16) -> u16 {
        let mut sym = table[(self.bit_buf >> (BITBUFSIZ - index_bits)) as usize];
        if sym >= threshold {
            let mut mask = 1u32 << (BITBUFSIZ - 1 - index_bits);
            loop {
                sym = if self.bit_buf & mask != 0 {
                    self.right[sym as usize]
                } else {
                    self.left[sym as usize]
                };
                mask >>= 1;
                if sym < threshold {
                    break;
                }
            }
        }
        sym
    }

    // -----------------------------------------------------------------------
    // Block header readers
    // -----------------------------------------------------------------------

    /// Read PT code lengths (`mPTLen`) and build `mPTTable`.
    ///
    /// Mirrors `ReadPTLen()` in Decompress.c.
    ///
    /// - `nn`: number of symbols (NT or MAXNP depending on call site)
    /// - `nbit`: bits used to encode the symbol count
    /// - `special`: the symbol index that triggers the run-length skip
    fn read_pt_len(&mut self, nn: u16, nbit: u16, special: u16) -> u16 {
        let number = self.get_bits(nbit) as u16;

        if number == 0 {
            let ch = self.get_bits(nbit) as u16;
            // Fill entire PT table with the single code, and zero all lengths.
            self.pt_table.fill(ch);
            self.pt_len[..nn as usize].fill(0);
            return 0;
        }

        let mut idx = 0usize;
        while idx < number as usize && idx < NPT {
            // Read next code length (3-bit with extension for values >= 7).
            let mut ch = (self.bit_buf >> (BITBUFSIZ - 3)) as u16;
            if ch == 7 {
                let mut mask = 1u32 << (BITBUFSIZ - 1 - 3);
                while self.bit_buf & mask != 0 {
                    mask >>= 1;
                    ch += 1;
                }
            }
            self.fill_buf(if ch < 7 { 3 } else { ch - 3 });
            self.pt_len[idx] = ch as u8;
            idx += 1;

            // Special run-length encoding at the `special` position.
            if idx as u16 == special {
                let run = (self.get_bits(2) as usize).min(NPT - idx);
                self.pt_len[idx..idx + run].fill(0);
                idx += run;
            }
        }

        // Zero-pad remaining entries up to nn.
        let end = (nn as usize).min(NPT);
        if idx < end {
            self.pt_len[idx..end].fill(0);
        }

        // Clone to satisfy borrow checker: make_table needs both pt_len and left/right/pt_table.
        let pt_len_copy: [u8; NPT] = *self.pt_len;
        Self::make_table(
            &mut self.left,
            &mut self.right,
            nn,
            &pt_len_copy[..nn as usize],
            8,
            self.pt_table.as_mut(),
        )
    }

    /// Read C code lengths (`mCLen`) and build `mCTable`.
    ///
    /// Mirrors `ReadCLen()` in Decompress.c.
    fn read_c_len(&mut self) {
        let number = self.get_bits(CBIT) as u16;

        if number == 0 {
            let ch = self.get_bits(CBIT) as u16;
            self.c_len.fill(0);
            self.c_table.fill(ch);
            return;
        }

        let mut idx = 0usize;
        while idx < number as usize {
            if idx >= NC {
                self.bad_table_flag = BAD_TABLE;
                return;
            }

            // Decode the next PT symbol (top 8 bits), then consume its code bits.
            let ch = self.decode_symbol(self.pt_table.as_ref(), 8, NT as u16);
            self.fill_buf(self.pt_len[ch as usize] as u16);

            if ch <= 2 {
                // Run-length codes.
                let count: i32 = if ch == 0 {
                    1
                } else if ch == 1 {
                    self.get_bits(4) as i32 + 3
                } else {
                    // ch == 2
                    self.get_bits(CBIT) as i32 + 20
                };
                let run = (count as usize).min(NC - idx);
                self.c_len[idx..idx + run].fill(0);
                idx += run;
            } else {
                self.c_len[idx] = (ch - 2) as u8;
                idx += 1;
            }
        }

        // Zero-pad remaining C lengths.
        self.c_len[idx..].fill(0);

        // Build the C table. Clone c_len to satisfy the borrow checker.
        let c_len_copy: [u8; NC] = *self.c_len;
        let result = Self::make_table(
            &mut self.left,
            &mut self.right,
            NC as u16,
            &c_len_copy,
            12,
            self.c_table.as_mut(),
        );
        if result == BAD_TABLE {
            self.bad_table_flag = BAD_TABLE;
        }
    }

    // -----------------------------------------------------------------------
    // Symbol decoders
    // -----------------------------------------------------------------------

    /// Decode one char/length code.
    ///
    /// Mirrors `DecodeC()` in Decompress.c. Starts a new block when `block_size == 0`.
    ///
    /// Returns `Err` when a block-size underflow or bad table is detected.
    fn decode_c(&mut self) -> Result<u16, DecompressError> {
        if self.block_size == 0 {
            // Start a new block: read block size, then all three tables.
            self.block_size = self.get_bits(16) as u16;

            self.bad_table_flag = self.read_pt_len(NT as u16, TBIT, 3);
            if self.bad_table_flag == 0 {
                self.read_c_len();
            }
            if self.bad_table_flag == 0 {
                // MAXNP is the threshold used by the C source (not NPT).
                self.bad_table_flag = self.read_pt_len(MAXNP as u16, self.pbit, u16::MAX);
            }
            if self.bad_table_flag != 0 {
                return Ok(0);
            }
        }

        if self.block_size == 0 {
            return Err(DecompressError::InvalidData("block size underflow"));
        }
        self.block_size -= 1;

        // Decode the next C symbol (top 12 bits), then consume its code bits.
        let idx = self.decode_symbol(self.c_table.as_ref(), 12, NC as u16);
        self.fill_buf(self.c_len[idx as usize] as u16);
        Ok(idx)
    }

    /// Decode one position value.
    ///
    /// Mirrors `DecodeP()` in Decompress.c.
    fn decode_p(&mut self) -> u32 {
        // Decode the next P symbol (top 8 bits), then consume its code bits.
        let val = self.decode_symbol(self.pt_table.as_ref(), 8, MAXNP as u16);
        self.fill_buf(self.pt_len[val as usize] as u16);

        let pos = val as u32;
        if val > 1 {
            // Low bits of the position.
            (1u32 << (val - 1)) | self.get_bits(val - 1)
        } else {
            pos
        }
    }

    // -----------------------------------------------------------------------
    // Main decode loop
    // -----------------------------------------------------------------------

    /// Decompress and return the output buffer.
    fn decode(mut self) -> Result<Vec<u8>, DecompressError> {
        // Prime the bit buffer.
        self.fill_buf(BITBUFSIZ);

        while self.out_buf < self.orig_size {
            let ch = self.decode_c()?;
            if self.bad_table_flag != 0 {
                return Err(DecompressError::InvalidData("bad Huffman table"));
            }

            if ch < 256 {
                // Literal byte.
                self.dst[self.out_buf as usize] = ch as u8;
                self.out_buf += 1;
            } else {
                // Back-reference: length encoded as ch - (0xFF + 1 - THRESHOLD) = ch - 253.
                let match_len = ch as u32 - (0xFF_u32 + 1 - THRESHOLD); // ch - 253

                let back = self.decode_p();
                // decode_p() does not set bad_table_flag; PT table errors are
                // caught earlier in decode_c() when read_pt_len() returns BAD_TABLE.

                // data_idx is a wrapping offset back into the sliding window.
                let data_idx = self.out_buf.wrapping_sub(back).wrapping_sub(1);

                for i in 0..match_len {
                    if self.out_buf >= self.orig_size {
                        break;
                    }
                    let src_idx = data_idx.wrapping_add(i) as usize;
                    if src_idx >= self.orig_size as usize {
                        return Err(DecompressError::InvalidData("match position out of range"));
                    }
                    self.dst[self.out_buf as usize] = self.dst[src_idx];
                    self.out_buf += 1;
                }
            }
        }

        Ok(self.dst)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // Build a minimal valid header (compressed_size, orig_size).
    fn make_header(comp_size: u32, orig_size: u32) -> Vec<u8> {
        let mut v = Vec::with_capacity(8);
        v.extend_from_slice(&comp_size.to_le_bytes());
        v.extend_from_slice(&orig_size.to_le_bytes());
        v
    }

    // -----------------------------------------------------------------------
    // parse_header tests
    // -----------------------------------------------------------------------

    #[test]
    fn parse_header_too_short() {
        assert_eq!(parse_header(&[]), Err(DecompressError::InputTooShort));
        assert_eq!(parse_header(&[0u8; 7]), Err(DecompressError::InputTooShort));
    }

    #[test]
    fn parse_header_exact_min() {
        // 8-byte header with comp_size=0: payload is empty, total = 8 bytes.
        let hdr = make_header(0, 0);
        let h = parse_header(&hdr).unwrap();
        assert_eq!(h.original_size, 0);
        assert_eq!(h.compressed_size, 0);
    }

    #[test]
    fn parse_header_comp_size_exceeds_src() {
        // comp_size claims 100 bytes but we only provide 8.
        let hdr = make_header(100, 200);
        assert_eq!(
            parse_header(&hdr),
            Err(DecompressError::InvalidData("compressed size exceeds source buffer"))
        );
    }

    #[test]
    fn parse_header_valid_with_payload() {
        let mut data = make_header(4, 10);
        data.extend_from_slice(&[0xAA, 0xBB, 0xCC, 0xDD]);
        let h = parse_header(&data).unwrap();
        assert_eq!(h.original_size, 10);
        assert_eq!(h.compressed_size, 4);
    }

    // -----------------------------------------------------------------------
    // Zero original size
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_zero_orig_size_returns_empty() {
        // comp_size=0, orig_size=0 → short-circuits before decoding, returns Ok([]).
        let hdr = make_header(0, 0);
        let out = decompress(&hdr, CompressionType::EfiStandard).unwrap();
        assert!(out.is_empty());
    }

    // -----------------------------------------------------------------------
    // block_size = 0 regression (must not panic, must return Err)
    //
    // When get_bits(16) reads a block-size field of 0x0000 from the stream,
    // block_size is still 0 after the block-init branch.  In the C original
    // this wraps (u16 underflow); in Rust debug mode it panics.  The fix must
    // return DecompressError::InvalidData instead.
    //
    // Trigger: orig_size > 0 so the decode loop runs at least once, and the
    // first 16 bits of the payload are 0x0000 (block-size field = 0).
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_block_size_zero_returns_error() {
        // orig_size = 4 forces the decode loop to run.
        // Payload starts with 0x00 0x00 … so get_bits(16) == 0, leaving block_size == 0.
        let mut data = make_header(16, 4);
        data.extend_from_slice(&[0x00_u8; 16]);
        let result = decompress(&data, CompressionType::EfiStandard);
        assert!(
            result.is_err(),
            "expected Err for block_size=0 stream, got Ok"
        );
        assert_eq!(
            result.unwrap_err(),
            DecompressError::InvalidData("block size underflow"),
        );
    }

    // -----------------------------------------------------------------------
    // Known test vectors — generated by uefi_firmware.efi_compressor.EfiCompress
    // and verified with EfiDecompress (Python C extension wrapping the EDK2
    // reference implementation).
    // -----------------------------------------------------------------------

    /// EFI-compressed encoding of the single ASCII byte b'A' (0x41).
    ///
    /// Generated with: `EfiCompress(b'A', 1)`
    /// Header: comp_size=8, orig_size=1
    fn efi_compressed_single_a() -> Vec<u8> {
        vec![
            0x08, 0x00, 0x00, 0x00, // comp_size = 8
            0x01, 0x00, 0x00, 0x00, // orig_size = 1
            0x00, 0x01, 0x00, 0x00, 0x04, 0x10, 0x00, 0x00,
        ]
    }

    /// EFI-compressed encoding of the single zero byte 0x00.
    ///
    /// Generated with: `EfiCompress(b'\x00', 1)`
    /// Header: comp_size=8, orig_size=1
    fn efi_compressed_single_zero() -> Vec<u8> {
        vec![
            0x08, 0x00, 0x00, 0x00, // comp_size = 8
            0x01, 0x00, 0x00, 0x00, // orig_size = 1
            0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        ]
    }

    /// EFI-compressed encoding of b"Hello, World!" (13 bytes).
    ///
    /// Generated with: `EfiCompress(b'Hello, World!', 13)`
    fn efi_compressed_hello_world() -> Vec<u8> {
        vec![
            0x18, 0x00, 0x00, 0x00, // comp_size = 24
            0x0d, 0x00, 0x00, 0x00, // orig_size = 13
            0x00, 0x0d, 0x3a, 0x51, 0x8d, 0x1c, 0xf8, 0x19, 0x4b, 0xdc, 0x07, 0x9b, 0x99, 0xd9,
            0x3f, 0x0c, 0x10, 0x06, 0xb0, 0x99, 0x5d, 0x3c, 0x56, 0x00,
        ]
    }

    /// EFI-compressed encoding of 32 repetitions of b'A'.
    ///
    /// Generated with: `EfiCompress(b'A' * 32, 32)`
    fn efi_compressed_repeated_a() -> Vec<u8> {
        vec![
            0x0b, 0x00, 0x00, 0x00, // comp_size = 11
            0x20, 0x00, 0x00, 0x00, // orig_size = 32
            0x00, 0x02, 0x20, 0x04, 0x31, 0xd0, 0xb6, 0x63, 0x40, 0x10, 0x00,
        ]
    }

    #[test]
    fn parse_header_known_vector() {
        let data = efi_compressed_single_zero();
        let h = parse_header(&data).unwrap();
        assert_eq!(h.original_size, 1);
        assert_eq!(h.compressed_size, 8);
    }

    #[test]
    fn decompress_single_byte_a() {
        let data = efi_compressed_single_a();
        let out = decompress(&data, CompressionType::EfiStandard).unwrap();
        assert_eq!(out, b"A");
    }

    #[test]
    fn decompress_single_byte_zero() {
        let data = efi_compressed_single_zero();
        let out = decompress(&data, CompressionType::EfiStandard).unwrap();
        assert_eq!(out, &[0x00]);
    }

    #[test]
    fn decompress_hello_world() {
        let data = efi_compressed_hello_world();
        let out = decompress(&data, CompressionType::EfiStandard).unwrap();
        assert_eq!(out, b"Hello, World!");
    }

    #[test]
    fn decompress_repeated_byte_backreference() {
        // 32 × b'A' exercises the back-reference (copy) path.
        let data = efi_compressed_repeated_a();
        let out = decompress(&data, CompressionType::EfiStandard).unwrap();
        assert_eq!(out, vec![b'A'; 32]);
    }

    // -----------------------------------------------------------------------
    // Garbage input must return Err, not panic.
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_garbage_returns_error() {
        // A valid-length header with 0xFF garbage payload.
        let mut data = make_header(16, 4);
        data.extend_from_slice(&[0xFF_u8; 16]);
        let result = decompress(&data, CompressionType::EfiStandard);
        assert!(result.is_err(), "expected Err for garbage input, got Ok");
    }

    #[test]
    fn decompress_garbage_tiano_returns_error() {
        let mut data = make_header(16, 4);
        data.extend_from_slice(&[0xFF_u8; 16]);
        let result = decompress(&data, CompressionType::Tiano);
        assert!(result.is_err(), "expected Err for garbage Tiano input, got Ok");
    }

    // -----------------------------------------------------------------------
    // Compression type enum
    // -----------------------------------------------------------------------

    #[test]
    fn compression_type_equality() {
        assert_eq!(CompressionType::EfiStandard, CompressionType::EfiStandard);
        assert_eq!(CompressionType::Tiano, CompressionType::Tiano);
        assert_ne!(CompressionType::EfiStandard, CompressionType::Tiano);
    }

    // -----------------------------------------------------------------------
    // Error display
    // -----------------------------------------------------------------------

    #[test]
    fn error_display() {
        let e = DecompressError::InputTooShort;
        assert_eq!(format!("{e}"), "input too short");

        let e = DecompressError::InvalidData("test");
        assert_eq!(format!("{e}"), "invalid data: test");
    }

    #[test]
    fn error_is_std_error() {
        fn check(_: &dyn std::error::Error) {}
        check(&DecompressError::InputTooShort);
        check(&DecompressError::InvalidData("x"));
    }

    // -----------------------------------------------------------------------
    // Fix 1: c_len index out-of-bounds guard — number == 511 (NC + 1)
    //
    // A malformed block claims 511 symbols in the C-length table (NC = 510).
    // The loop must detect idx >= NC and return InvalidData, not panic.
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_c_len_number_511_returns_error() {
        // Build a payload that makes get_bits(CBIT=9) return 511 (0x1FF).
        // The first 16 bits of the payload are the block-size field; after
        // fill_buf(32) primes the bit-buffer with the first 4 bytes, the
        // decoder reads:
        //   - 16 bits → block_size
        //   - 5 bits  → PT count for read_pt_len(NT=19, TBIT=5, special=3)
        // We arrange the bytes so that read_pt_len returns a single-code PT
        // table (count == 0 path) and then read_c_len sees number == 511.
        // Easiest: fill with 0xFF so every multi-bit read returns all-ones.
        // With all-ones:
        //   block_size  = 0xFFFF (large, fine)
        //   PT count (5 bits) = 31 ≠ 0, so we go into the loop path, which
        //   may not build a valid PT table, triggering bad_table_flag.
        // Use 0x00 instead: all zero-bits means every count field == 0,
        // selecting the single-code shortcut for every table:
        //   block_size = 0x0000 → triggers "block size underflow" before
        //   we even reach read_c_len.
        //
        // Construct a targeted byte sequence:
        //   bytes 0-1 (block_size 16 bits): 0x00 0x01  → block_size = 256
        //   next 5 bits (PT count=TBIT):    0b00000    → count == 0 → single code PT
        //   next 5 bits (PT ch value):      0b00000    → ch == 0
        //   next 9 bits (C count=CBIT):     0b111111111 → number == 511
        // Pack into a byte stream (MSB-first bit packing):
        //   byte 0: 0x00
        //   byte 1: 0x01
        //   byte 2...: 0b00000_000 | 0b00_0000111 | 0b11111_xxx
        // Full bit sequence: 00000000 00000001 00000000 00001111 11111???
        //   b0=0x00, b1=0x01, b2=0x00, b3=0x0F, b4=0xFF, remainder don't care
        let mut data = make_header(32, 4);
        data.push(0x00); // block_size high byte (LE bit-stream MSB first)
        data.push(0x01); // block_size = 256
        data.push(0x00); // PT count (5 bits = 0) + PT ch bits start
        data.push(0x0F); // remaining PT ch bits + C count bits start
        data.push(0xFF); // C count bits continue: should yield 511
        data.resize(8 + 32, 0x00); // pad to declared comp_size
        let result = decompress(&data, CompressionType::EfiStandard);
        assert!(
            result.is_err(),
            "expected Err for c_len number=511 stream, got Ok"
        );
    }

    // -----------------------------------------------------------------------
    // Fix 2: oversized orig_size is rejected before allocation
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_oversized_orig_size_returns_error() {
        // Header claims to decompress to 65 MiB — well above the 64 MiB cap.
        let orig_size: u32 = 65 * 1024 * 1024;
        let comp_size: u32 = 8; // minimal payload
        let mut data = make_header(comp_size, orig_size);
        data.extend_from_slice(&[0u8; 8]);
        let result = decompress(&data, CompressionType::EfiStandard);
        assert_eq!(
            result,
            Err(DecompressError::InvalidData("original size exceeds maximum (64 MiB)"))
        );
    }

    // -----------------------------------------------------------------------
    // Fix 9: positive Tiano test vector
    //
    // Vector generated with Python:
    //   from uefi_firmware.efi_compressor import TianoCompress
    //   TianoCompress(b"Hello, Tiano!", 13)
    //   → [25, 0, 0, 0, 13, 0, 0, 0, 0, 13, 59, 78, 72, 225, 193, 148,
    //       59, 192, 241, 17, 40, 8, 22, 206, 144, 3, 19, 112, 86, 241,
    //       249, 96, 0]
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_tiano_known_vector() {
        let compressed: &[u8] = &[
            25, 0, 0, 0, 13, 0, 0, 0, // header: comp=25, orig=13
            0, 13, 59, 78, 72, 225, 193, 148, 59, 192, 241, 17, 40, 8, 22, 206, 144, 3, 19, 112,
            86, 241, 249, 96, 0,
        ];
        let result = decompress(compressed, CompressionType::Tiano).unwrap();
        assert_eq!(result, b"Hello, Tiano!");
    }

    // -----------------------------------------------------------------------
    // Fix 10: truncated stream (header-only, no payload data) returns error
    // -----------------------------------------------------------------------

    #[test]
    fn decompress_truncated_stream_returns_error() {
        // Valid header claiming 100 bytes original, but only 8 bytes total (header only).
        let mut src = vec![0u8; 8];
        src[0..4].copy_from_slice(&8u32.to_le_bytes()); // compressed_size = 8
        src[4..8].copy_from_slice(&100u32.to_le_bytes()); // original_size = 100
        let result = decompress(&src, CompressionType::EfiStandard);
        assert!(result.is_err());
    }
}
