Going-Commando/crates/gc-formats/src/wad.rs

//! Insomniac's "WAD" LZ77 container.
//!
//! Spec: `docs/formats/wad.md`.
//!
//! A 16-byte header followed by a stream of packets. Each packet begins with a
//! **flag byte** that selects one of four packet classes:
//!
//! | Flag range | Packet |
//! |---|---|
//! | `0x00..0x0f` | literal run (the flag byte is the count, biased) |
//! | `0x10..0x1f` | back-reference, 15-bit displacement, short length |
//! | `0x20..0x3f` | back-reference, 14-bit displacement |
//! | `0x40..0xff` | back-reference, 11-bit displacement, short length |
//!
//! Every packet ends with 0-3 extra literal bytes carried in unused bits, so
//! literal and back-reference packets interleave without alignment. Two
//! special cases exist in the `0x10..0x1f` class: a packet with a zero far
//! field is either a **pad packet** (skip forward to the next 0x1000 boundary)
//! or a **dummy packet**, and zero displacements are invalid.
//!
//! The encoding is verified against retail data: every WAD on the `SCUS_972.68`
//! disc decodes, and each decompresses to a size its containing header agrees
//! with.

use crate::buf::{Error, ErrorKind, Result};

pub const WAD_HEADER_SIZE: usize = 0x10;

/// Is this buffer a WAD stream?
pub fn is_wad(b: &[u8]) -> bool {
    b.len() >= WAD_HEADER_SIZE && &b[..3] == b"WAD"
}

/// Total stream length in bytes, including the 16-byte header.
pub fn compressed_size(b: &[u8]) -> Result<u32> {
    if !is_wad(b) {
        return Err(Error {
            offset: 0,
            context: "WAD header",
            kind: ErrorKind::BadMagic { expected: "WAD" },
        });
    }
    Ok(u32::from_le_bytes([b[3], b[4], b[5], b[6]]))
}

/// Decompress a WAD stream.
pub fn decompress(bytes: &[u8]) -> Result<Vec<u8>> {
    let total = compressed_size(bytes)? as usize;
    if total < WAD_HEADER_SIZE || total > bytes.len() {
        return Err(Error {
            offset: 0,
            context: "WAD header",
            kind: ErrorKind::InvalidLength("WAD stream size"),
        });
    }

    let stream = &bytes[WAD_HEADER_SIZE..total];
    let end = stream.len();
    let mut ptr = 0usize;
    let mut out: Vec<u8> = Vec::with_capacity(total * 2);

    while ptr < end {
        let flag_at = ptr;
        let flag = stream[ptr];
        ptr += 1;

        if flag < 0x10 {
            // Literal run. 0x00 is the escape for a count stored in the next byte.
            let n = if flag == 0 {
                let extra = take(stream, &mut ptr, "WAD literal length")?;
                extra as usize + 18
            } else {
                flag as usize + 3
            };
            take_into(stream, &mut ptr, &mut out, n, "WAD literal run")?;
            // Two literal packets in a row are not a valid encoding: the second
            // would have been merged into the first.
            if ptr < end && stream[ptr] < 0x10 {
                return Err(Error {
                    offset: WAD_HEADER_SIZE + ptr,
                    context: "WAD literal run",
                    kind: ErrorKind::Invalid("adjacent literal packets"),
                });
            }
            continue;
        }

        let (match_len, displacement, little): (usize, usize, u8);
        if flag < 0x20 {
            // 15-bit displacement window.
            let mut ml = (flag & 7) as usize;
            if ml == 0 {
                ml = take(stream, &mut ptr, "WAD match length")? as usize + 7;
            }
            let b0 = take(stream, &mut ptr, "WAD back-reference")?;
            let b1 = take(stream, &mut ptr, "WAD back-reference")?;
            let a = ((flag >> 3) & 1) as usize;
            let far = b1 as usize * 0x40 + (b0 >> 2) as usize;

            if a == 0 && far == 0 {
                if ml != 1 {
                    // Pad packet: skip to the next 0x1000 boundary.
                    ptr = (ptr + 0xFFF) & !0xFFF;
                    continue;
                }
                // Dummy packet: carries only its trailing literals.
                take_into(
                    stream,
                    &mut ptr,
                    &mut out,
                    (b0 & 3) as usize,
                    "WAD dummy packet",
                )?;
                continue;
            }

            displacement = 0x4000 * (a + 1) + far;
            match_len = ml + 2;
            little = b0;
        } else if flag < 0x40 {
            // 14-bit displacement.
            let mut ml = (flag & 0x1f) as usize;
            if ml == 0 {
                ml = take(stream, &mut ptr, "WAD match length")? as usize + 0x1f;
            }
            match_len = ml + 2;
            let b1 = take(stream, &mut ptr, "WAD back-reference")?;
            let b2 = take(stream, &mut ptr, "WAD back-reference")?;
            displacement = b2 as usize * 0x40 + (b1 >> 2) as usize + 1;
            little = b1;
        } else {
            // 11-bit displacement, short length.
            let b1 = take(stream, &mut ptr, "WAD back-reference")?;
            match_len = (flag >> 5) as usize + 1;
            displacement = b1 as usize * 8 + ((flag >> 2) & 7) as usize + 1;
            little = flag;
        }

        copy_match(&mut out, displacement, match_len, WAD_HEADER_SIZE + flag_at)?;
        take_into(
            stream,
            &mut ptr,
            &mut out,
            (little & 3) as usize,
            "WAD trailing literals",
        )?;
    }

    Ok(out)
}

/// Read one byte, advancing the cursor.
fn take(stream: &[u8], ptr: &mut usize, context: &'static str) -> Result<u8> {
    if *ptr >= stream.len() {
        return Err(Error {
            offset: WAD_HEADER_SIZE + *ptr,
            context,
            kind: ErrorKind::OutOfBounds {
                needed: 1,
                available: stream.len().saturating_sub(*ptr),
            },
        });
    }
    let v = stream[*ptr];
    *ptr += 1;
    Ok(v)
}

/// Copy `n` literal bytes from the stream into `out`, advancing the cursor.
fn take_into(
    stream: &[u8],
    ptr: &mut usize,
    out: &mut Vec<u8>,
    n: usize,
    context: &'static str,
) -> Result<()> {
    if stream.len() - *ptr < n {
        return Err(Error {
            offset: WAD_HEADER_SIZE + *ptr,
            context,
            kind: ErrorKind::OutOfBounds {
                needed: n,
                available: stream.len().saturating_sub(*ptr),
            },
        });
    }
    out.extend_from_slice(&stream[*ptr..*ptr + n]);
    *ptr += n;
    Ok(())
}

/// Copy `len` bytes from `displacement` back in `out`.
///
/// Overlapping copies expand as run-length: the source index advances with the
/// destination, so displacement 1 repeats a single byte.
fn copy_match(out: &mut Vec<u8>, displacement: usize, len: usize, at: usize) -> Result<()> {
    if displacement == 0 || displacement > out.len() {
        return Err(Error {
            offset: at,
            context: "WAD back-reference",
            kind: ErrorKind::Invalid("match source before the start of the output"),
        });
    }
    let src = out.len() - displacement;
    for i in 0..len {
        let v = out[src + i];
        out.push(v);
    }
    Ok(())
}

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

    fn header(total: u32) -> Vec<u8> {
        let mut h = vec![0u8; WAD_HEADER_SIZE];
        h[..3].copy_from_slice(b"WAD");
        h[3..7].copy_from_slice(&total.to_le_bytes());
        h
    }

    fn stream(body: Vec<u8>) -> Vec<u8> {
        let mut s = header((body.len() + WAD_HEADER_SIZE) as u32);
        s.extend_from_slice(&body);
        s
    }

    #[test]
    fn recognises_the_magic() {
        assert!(is_wad(&header(0x10)));
        assert!(!is_wad(b"NOT"));
        assert!(!is_wad(b"WAD"));
    }

    #[test]
    fn rejects_a_non_wad() {
        let err = decompress(&[0u8; 32]).unwrap_err();
        assert!(matches!(err.kind, ErrorKind::BadMagic { .. }));
    }

    #[test]
    fn rejects_a_size_beyond_the_buffer() {
        let mut s = header(0x1000);
        s.extend(std::iter::repeat_n(0u8, 8));
        assert!(matches!(
            decompress(&s).unwrap_err().kind,
            ErrorKind::InvalidLength(_)
        ));
    }

    #[test]
    fn rejects_a_size_below_the_header() {
        assert!(decompress(&header(4)).is_err());
    }

    #[test]
    fn literal_run() {
        // Flag 0x06 -> 6 + 3 = 9 literal bytes.
        let mut body = vec![0x06];
        body.extend(b"abcdefghi");
        assert_eq!(decompress(&stream(body)).unwrap(), b"abcdefghi");
    }

    #[test]
    fn long_literal_run_via_escape() {
        // Flag 0x00 then a length byte: 20 + 18 = 38 literals.
        let payload: Vec<u8> = (0..38u8).collect();
        let mut body = vec![0x00, 20u8];
        body.extend_from_slice(&payload);
        assert_eq!(decompress(&stream(body)).unwrap(), payload);
    }

    #[test]
    fn eleven_bit_back_reference() {
        // Seed 9 literals so a displacement is available.
        let mut body = vec![0x06]; // 9 literals
        body.extend(b"abcdefghi");
        // Flag 0x40..: match_len = (flag>>5)+1, disp = b1*8 + ((flag>>2)&7) + 1,
        // trailing literals = flag & 3.
        // flag 0x40 -> match_len = (0x40 >> 5) + 1 = 3.
        // displacement = b1 * 8 + ((flag >> 2) & 7) + 1 = 1 * 8 + 0 + 1 = 9.
        body.push(0x40);
        body.push(0x01);
        assert_eq!(&decompress(&stream(body)).unwrap()[..12], b"abcdefghiabc");
    }

    #[test]
    fn overlapping_back_reference_is_run_length() {
        // displacement 1 expands to a repeated byte.
        let mut body = vec![0x03]; // 6 literals
        body.extend(b"zzzzzz");
        // flag 0x40, b1 = 0 -> disp = 0*8 + 0 + 1 = 1, len = (0x40>>5)+1 = 3.
        body.push(0x40);
        body.push(0x00);
        assert_eq!(&decompress(&stream(body)).unwrap()[..9], b"zzzzzzzzz");
    }

    #[test]
    fn rejects_a_back_reference_before_the_start() {
        let body = vec![0x40, 0x00];
        assert!(decompress(&stream(body)).is_err());
    }

    #[test]
    fn dummy_packet_carries_only_literals() {
        // 0x10 class with a = 0, far = 0 and ml == 1 is a dummy packet:
        // flag = 0x10 | 1 = 0x11, then b0, b1 = 0x00, 0x00 -> b0 & 3 = 0 literals.
        let mut body = vec![0x06];
        body.extend(b"abcdefghi");
        body.extend_from_slice(&[0x11, 0x00, 0x00]);
        let out = decompress(&stream(body)).unwrap();
        assert_eq!(&out[..9], b"abcdefghi");
    }

    #[test]
    fn adjacent_literal_packets_are_rejected() {
        let mut body = vec![0x03];
        body.extend(b"abcd");
        body.push(0x03);
        body.extend(b"efgh");
        assert!(decompress(&stream(body)).is_err());
    }

    #[test]
    fn empty_body_decompresses_to_nothing() {
        let s = header(WAD_HEADER_SIZE as u32);
        assert!(decompress(&s).unwrap().is_empty());
    }
}