Going-Commando/tools/trace/src/tfrag_cmp.rs

//! The tfrag lighting comparison: the game's own bytes against ours.
//!
//! `LightTfrags` overwrites every tfrag's RGBA block in EE RAM with its lit
//! per-vertex colours, so after a level loads the RAM holds the answer. The
//! comparison is:
//!
//! 1. **Locate** the tfrags block in RAM by content, never by a fixed
//!    address — the load address differs per level and per run. The tfrag-0
//!    bounding sphere is the search key; a candidate is only accepted when
//!    *every* header matches the disc on the bytes the game leaves alone
//!    (all but the relocated `data` pointer and the runtime occlusion
//!    index) and the `data` pointers are relocated exactly as
//!    `table + disc_data`.
//! 2. **Locate the light bank** the same way: the bytes of the level's
//!    gameplay light section, searched in RAM. It is compared against the
//!    disc's so a parse error and a maths error are distinguishable.
//! 3. **Run the pass** over the same disc data with
//!    [`gc_engine::tfrag_light::light_tfrag`].
//! 4. **Diff** per channel, per tfrag, bucketed by vertex class so a
//!    mismatch in the blend path (which is the only path `rsqrt` touches)
//!    separates from one in the plain path.
//!
//! The pass-ran check is the number of vertices whose RAM colour still equals
//! the disc's unlit placeholder. If that is near the whole level, the pass
//! has not run yet and the comparison means nothing.

use std::collections::BTreeMap;
use std::path::Path;

use anyhow::{anyhow, Context, Result};

use gc_engine::tfrag_light::{light_tfrag, light_vertex_terms, vertex_normal_bits, VertexTerms};
use gc_formats::level_core::LevelCoreHeader;
use gc_formats::tfrag::{parse_tfrag, TfragHeader, TfragsBlock, TFRAG_HEADER_SIZE};
use gc_formats::tfrag_light::{LightBank, NormalTable};

use crate::ee::EeRam;

/// How many mismatching vertices keep their terms in the report. Enough to
/// read a pattern by hand, few enough that the report stays readable.
const BREAKDOWNS: usize = 6;

/// Result of one comparison, everything a report needs.
#[derive(Debug, Clone)]
pub struct Report {
    pub level: usize,
    /// RAM address of the tfrags block, its header table, and the tfrag count.
    pub block_base: usize,
    pub table: usize,
    pub tfrag_count: usize,
    pub headers_matched: usize,
    pub pointers_relocated: usize,
    /// RAM address of the directional light bank, when found.
    pub bank_ram: Option<usize>,
    /// Whether the bank was found in RAM **byte-identical to the disc's**. It
    /// is a content hit, so a found bank is by construction the disc's; the
    /// flag exists to keep the not-found case out of that claim, because then
    /// the comparison below runs on the disc's bank with nothing from the
    /// game to confirm it.
    pub bank_is_disc: bool,
    pub vertices_compared: usize,
    pub bytes_compared: usize,
    pub bytes_equal: usize,
    /// Vertices whose RAM colour still equals the disc's unlit placeholder.
    pub placeholder_equal: usize,
    /// Per-vertex-class tallies: `single-set`, `blended`, and `no-set` for
    /// vertices whose `select` names a light set the disc's bank does not have
    /// -- runtime lights, which a level-load comparison cannot see.
    pub classes: BTreeMap<&'static str, ClassStat>,
    /// `ours - ram` per channel, in order.
    pub histogram: BTreeMap<i32, usize>,
    /// The first few differing vertices, for reading by hand.
    pub mismatches: Vec<Mismatch>,
    /// The pass's terms for the first few mismatching vertices. A colour
    /// mismatch alone cannot be attributed, and the whole point of the
    /// comparison is to say *which* term the game disagrees with.
    pub breakdowns: Vec<Breakdown>,
}

/// One mismatching vertex with its terms, so the difference can be read as
/// `ours` versus `ram` per term rather than per byte.
#[derive(Debug, Clone)]
pub struct Breakdown {
    pub tfrag: usize,
    pub entry: usize,
    /// `None` when the pass found no light set for the vertex's `select` --
    /// which is itself the finding, so it is recorded rather than skipped.
    pub terms: Option<VertexTerms>,
    /// The record's `select`, always available even when the terms are not.
    pub select: u16,
    /// The record's base colour, PEXT5-expanded: what the game would write if
    /// the light contributed nothing.
    pub base: [u8; 4],
    /// The disc's stored RGBA for the entry, which is the unlit placeholder.
    pub disc: [u8; 4],
    pub ram: [u8; 4],
}

/// Totals for one vertex class.
#[derive(Debug, Clone, Copy, Default)]
pub struct ClassStat {
    pub vertices: usize,
    pub equal: usize,
}

/// One differing vertex byte.
#[derive(Debug, Clone, Copy)]
pub struct Mismatch {
    pub tfrag: usize,
    pub entry: usize,
    pub channel: usize,
    pub ours: i32,
    pub ram: i32,
}

impl Report {
    /// 0 = every byte equal, 2 = mismatches.
    pub fn exit_code(&self) -> i32 {
        if self.bytes_equal == self.bytes_compared {
            0
        } else {
            2
        }
    }
}

/// Everything the comparison needs from one unpacked level.
pub struct LevelInputs {
    pub core_data: Vec<u8>,
    pub core_index: Vec<u8>,
    pub gameplay: Vec<u8>,
    pub overlay: Vec<u8>,
}

impl LevelInputs {
    pub fn load(dir: &Path, level: usize) -> Result<LevelInputs> {
        let read = |name: &str| {
            std::fs::read(dir.join(format!("level{level}")).join(name))
                .with_context(|| format!("level{level}/{name}"))
        };
        Ok(LevelInputs {
            core_data: read("core_data.bin")?,
            core_index: read("core_index.bin")?,
            gameplay: read("gameplay.bin")?,
            overlay: read("overlay.bin")?,
        })
    }
}

/// Compare the tfrag lighting of one level against a RAM image.
pub fn compare(inputs: &LevelInputs, level: usize, ram: &EeRam) -> Result<Report> {
    // Parsed for its validation, not for a field: a core index that does not
    // parse cannot name the tfrags region the rest of this reads.
    LevelCoreHeader::parse(&inputs.core_index)?;
    let tfrags_ofs = i32_at(&inputs.core_index, 0x08)
        .ok_or_else(|| anyhow!("core index is too short for the tfrags offset"))?;
    let region_end = [0x0c_usize, 0x10, 0x14]
        .iter()
        .filter_map(|&o| i32_at(&inputs.core_index, o))
        .filter(|&v| v > tfrags_ofs)
        .min()
        .unwrap_or(inputs.core_data.len() as i32);
    let region = &inputs.core_data[tfrags_ofs as usize..region_end as usize];
    let block = TfragsBlock::parse(region)?;
    let table_off = block.table_offset as usize;
    let count = block.tfrag_count as usize;
    if count == 0 {
        anyhow::bail!("level{level} has no tfrags");
    }

    let region_len = region.len();
    let block_base = locate_block(region, ram, table_off, count)?;

    // The RAM table, then every header against the disc on the invariant mask.
    let table = block_base + table_off;
    let mut headers_matched = 0usize;
    let mut pointers_relocated = 0usize;
    for i in 0..count {
        let at = i * TFRAG_HEADER_SIZE;
        let disc_h = &region[table_off + at..table_off + at + TFRAG_HEADER_SIZE];
        let ram_h = ram
            .slice(table + at, TFRAG_HEADER_SIZE)
            .ok_or_else(|| anyhow!("tfrag header {i} runs past the RAM image"))?;
        if invariant_equal(disc_h, ram_h) {
            headers_matched += 1;
        }
        // data is relocated to an absolute pointer into the RAM copy.
        let disc_data = i32_at(disc_h, 0x10).unwrap_or(0);
        let ram_data = ram.i32(table + at + 0x10).unwrap_or(0);
        if ram_data as usize == block_base + table_off + disc_data as usize {
            pointers_relocated += 1;
        }
    }

    // The bank: search RAM for the gameplay's light section, byte for byte.
    // What is searched for is the section's own bytes, so a hit *is* the
    // comparison -- a bank that is in RAM by content is the disc's bank, and
    // the only other outcome is that it is not there, which is reported
    // rather than papered over.
    let disc_bank = LightBank::parse(&inputs.gameplay)?;
    let section = bank_section(&inputs.gameplay)?;
    let bank_ram = if section.is_empty() {
        None
    } else {
        ram.find(section)
    };

    let table_normals =
        NormalTable::from_overlay(&inputs.overlay).context("the level overlay's normal table")?;

    let mut report = Report {
        level,
        block_base,
        table,
        tfrag_count: count,
        headers_matched,
        pointers_relocated,
        bank_ram,
        bank_is_disc: bank_ram.is_some(),
        vertices_compared: 0,
        bytes_compared: 0,
        bytes_equal: 0,
        placeholder_equal: 0,
        classes: BTreeMap::new(),
        histogram: BTreeMap::new(),
        mismatches: Vec::new(),
        breakdowns: Vec::new(),
    };

    for i in 0..count {
        let at = i * TFRAG_HEADER_SIZE;
        let disc_h =
            TfragHeader::parse(&region[table_off + at..table_off + at + TFRAG_HEADER_SIZE])?;
        let ram_h = TfragHeader::parse(ram.slice(table + at, TFRAG_HEADER_SIZE).unwrap())?;
        let vert = disc_h.vert_count as usize;
        if vert == 0 {
            continue;
        }
        // The tfrag's own data block in RAM, bounded by the region. `ram_h.data`
        // is the already-relocated absolute pointer, and the region ends at
        // `block_base + region_len`, so the block is what is between the two.
        let ram_data = ram_h.data as usize;
        let start = ram_data
            .checked_add(0)
            .filter(|&s| s < ram.bytes().len())
            .ok_or_else(|| anyhow!("tfrag {i} data pointer {ram_data:#x} is outside RAM"))?;
        let region_end_ram = block_base + region_len;
        if start >= region_end_ram {
            anyhow::bail!("tfrag {i} data pointer {ram_data:#x} is past the tfrags region");
        }
        let len = region_end_ram - start;
        let block_ram = &ram.bytes()[start..start + len];
        let geom = parse_tfrag(block_ram, &ram_h).context("the game's copy of tfrag {i}")?;

        // The lit colours, from our pass over this same data.
        let ours = light_tfrag(&geom, &disc_bank, &table_normals);

        // And the RAM's, plus the disc's unlit placeholder for the same
        // entries. Only 0..vert_count mean anything: past them the block is
        // scratchpad leftovers that no vertex references.
        let rgba_ofs = ram_h.rgba_ofs as usize;
        let entries = vert.min(ours.len());
        let disc_data = disc_h.data as usize;
        for (e, ours_byte) in ours.iter().take(entries).enumerate() {
            let l = &geom.lights[e];
            let select = gem_select(&geom, e);
            let class = if select & 0xff00 == 0 {
                if disc_bank.sets.len() > (select & 0xf) as usize {
                    "single-set"
                } else {
                    "runtime-set"
                }
            } else if disc_bank.sets.len() > (select & 0xf) as usize {
                "blended"
            } else {
                "runtime-blend"
            };
            let stat = report.classes.entry(class).or_default();
            stat.vertices += 1;
            report.vertices_compared += 1;
            let ram_rgba = match ram.slice(start + rgba_ofs + e * 4, 4) {
                Some(b) => b,
                None => break,
            };
            let disc_rgba = region
                .get(table_off + disc_data + disc_h.rgba_ofs as usize + e * 4..)
                .and_then(|b| b.get(..4))
                .unwrap_or(&[0, 0, 0, 0]);
            let mut all_equal = true;
            for (ch, ours_byte) in ours_byte.iter().enumerate() {
                let o = *ours_byte as i32;
                let r = ram_rgba[ch] as i32;
                report.bytes_compared += 1;
                if o == r {
                    report.bytes_equal += 1;
                } else {
                    all_equal = false;
                    *report.histogram.entry(o - r).or_default() += 1;
                    if report.mismatches.len() < 32 {
                        report.mismatches.push(Mismatch {
                            tfrag: i,
                            entry: e,
                            channel: ch,
                            ours: o,
                            ram: r,
                        });
                    }
                }
            }
            if all_equal {
                stat.equal += 1;
            } else if report.breakdowns.len() < BREAKDOWNS {
                // The terms for this vertex, as the pass computed them, so the
                // difference can be attributed to a term. A missing light set
                // is kept too: that is a finding, not a gap.
                let normal =
                    vertex_normal_bits(&table_normals, l.azimuth_index(), l.elevation_index());
                let terms = light_vertex_terms(&disc_bank, normal, &geom.lights[e], l.select);
                report.breakdowns.push(Breakdown {
                    tfrag: i,
                    entry: e,
                    terms,
                    select: l.select,
                    base: gc_engine::tfrag_light::pext5(l.color as u16),
                    disc: [disc_rgba[0], disc_rgba[1], disc_rgba[2], disc_rgba[3]],
                    ram: [ram_rgba[0], ram_rgba[1], ram_rgba[2], ram_rgba[3]],
                });
            }
            if ram_rgba == disc_rgba {
                report.placeholder_equal += 1;
            }
        }
    }

    Ok(report)
}

/// The `select` word for entry `e`, from the parsed light records.
fn gem_select(g: &gc_formats::tfrag::TfragGeometry, e: usize) -> u16 {
    g.lights.get(e).map(|l| l.select).unwrap_or(0)
}

/// The light bank's section bytes: the count word, the pad, then the sets.
///
/// This is what `LightBank::parse` reads, so searching RAM for it is checking
/// the parse against the game: a hit means the game holds the same bytes the
/// disc holds, and a miss means either the level rewrites the section or the
/// parse does not describe what the game loaded. Both are worth knowing
/// before any colour is compared.
fn bank_section(gameplay: &[u8]) -> Result<&[u8]> {
    let ofs = u32_at(gameplay, 0x04)
        .ok_or_else(|| anyhow!("gameplay extent is too short for the bank offset"))?
        as usize;
    let count = i32_at(gameplay, ofs)
        .ok_or_else(|| anyhow!("the light bank offset is past the gameplay extent"))?;
    if count < 0 {
        // No lights on this level, which is not an error.
        return Ok(&[]);
    }
    let end = ofs + 0x10 + 0x40 * count as usize;
    gameplay.get(ofs..end).ok_or_else(|| {
        anyhow!(
            "the light section runs {} bytes past the gameplay extent",
            end - gameplay.len()
        )
    })
}

/// Find the level's tfrags block in RAM.
///
/// The key is tfrag 0's bounding sphere: sixteen bytes of float that are
/// unique in a 32 MiB image, but the candidate still has to prove itself —
/// the block header must match, *every* header must match the disc on the
/// bytes the game does not rewrite, and every `data` pointer must be the
/// relocated form of the disc's. Anything less and the answer would be a
/// plausible but wrong read.
fn locate_block(region: &[u8], ram: &EeRam, table_off: usize, count: usize) -> Result<usize> {
    let key = &region[table_off..table_off + 16];
    let candidates = ram.find_all(key);
    let mut best: Option<(usize, usize)> = None;
    for &cand in &candidates {
        if cand < table_off {
            continue;
        }
        let base = cand - table_off;
        let Some(ram_block_header) = ram.slice(base, 0x10) else {
            continue;
        };
        if ram_block_header != &region[..0x10.min(region.len())] {
            continue;
        }
        let mut matched = 0usize;
        for i in 0..count.min(64) {
            let at = i * TFRAG_HEADER_SIZE;
            let disc_h = match region.get(table_off + at..table_off + at + TFRAG_HEADER_SIZE) {
                Some(h) => h,
                None => break,
            };
            match ram.slice(cand + at, TFRAG_HEADER_SIZE) {
                Some(h) if invariant_equal(disc_h, h) => matched += 1,
                _ => {}
            }
        }
        if best.is_none_or(|(_, m)| matched > m) {
            best = Some((base, matched));
        }
    }
    match best {
        Some((base, matched)) if matched == count.min(64) => Ok(base),
        Some((base, matched)) => Err(anyhow!(
            "tfrags block at {base:#x} matches {matched} of {count} headers on the invariant bytes; \
             refusing a partial match"
        )),
        None => Err(anyhow!(
            "no tfrags block in the RAM image matches this level ({} candidates for the tfrag-0 \
             bounding sphere)",
            candidates.len()
        )),
    }
}

/// Bytes a tfrag header must keep between the disc and RAM: everything except
/// the relocated `data` pointer (+0x10..0x14) and the occlusion index
/// (+0x3a..0x3c), which the loader renumbers per tfrag.
fn invariant_equal(disc: &[u8], ram: &[u8]) -> bool {
    let keep = |j: usize| !(0x10..0x14).contains(&j) && !(0x3a..0x3c).contains(&j);
    disc.iter()
        .zip(ram.iter())
        .enumerate()
        .filter(|&(j, _)| keep(j))
        .all(|(_, (a, b))| a == b)
}

fn i32_at(bytes: &[u8], at: usize) -> Option<i32> {
    bytes
        .get(at..at + 4)
        .map(|b| i32::from_le_bytes(b.try_into().unwrap()))
}

fn u32_at(bytes: &[u8], at: usize) -> Option<u32> {
    bytes
        .get(at..at + 4)
        .map(|b| u32::from_le_bytes(b.try_into().unwrap()))
}

/// Render a report for a terminal.
pub fn print_report(r: &Report, out: &mut dyn std::fmt::Write) -> std::fmt::Result {
    let _ = writeln!(
        out,
        "level{}: tfrags block at {:#x}, table at {:#x}, {} tfrags",
        r.level, r.block_base, r.table, r.tfrag_count
    );
    let _ = writeln!(
        out,
        "  headers matching the disc on invariant bytes: {}/{}",
        r.headers_matched, r.tfrag_count
    );
    let _ = writeln!(
        out,
        "  data pointers relocated as table + disc offset: {}/{}",
        r.pointers_relocated, r.tfrag_count
    );
    match r.bank_ram {
        Some(at) => writeln!(
            out,
            "  light bank in RAM at {at:#x}, byte-identical to the disc's"
        )?,
        None => writeln!(
            out,
            "  the disc's light bank is not in RAM by content; the colours below are\\
           \n    compared against the disc's bank with nothing from the game to confirm it"
        )?,
    }
    let _ = writeln!(
        out,
        "\nvertices compared: {} ({} colour bytes)",
        r.vertices_compared, r.bytes_compared
    );
    let _ = writeln!(
        out,
        "  equal:              {} ({:.2}%)",
        r.bytes_equal,
        100.0 * r.bytes_equal as f64 / r.bytes_compared.max(1) as f64
    );
    let _ = writeln!(
        out,
        "  unlit placeholders: {} of {} vertices ({:.2}%)",
        r.placeholder_equal,
        r.vertices_compared,
        100.0 * r.placeholder_equal as f64 / r.vertices_compared.max(1) as f64
    );
    for (class, stat) in &r.classes {
        let _ = writeln!(
            out,
            "  {class}: {} vertices, {} fully equal",
            stat.vertices, stat.equal
        );
    }
    let _ = writeln!(out, "\ndifference (ours - ram) per channel:");
    for d in -8..=8 {
        let n = r.histogram.get(&d).copied().unwrap_or(0);
        if n > 0 {
            let bar = "#".repeat((n as f64).log10().ceil() as usize);
            let _ = writeln!(out, "  {d:+3}: {n:8} {bar}");
        }
    }
    // The tail in magnitude buckets: a single far value hiding a count of
    // differing bytes is the number this report most needs to be legible.
    for (lo, hi) in [(9, 33), (33, 65), (65, 129), (129, 256), (256, i32::MAX)] {
        let n: usize = r
            .histogram
            .iter()
            .filter(|(&d, _)| !(-lo..=hi).contains(&d))
            .map(|(_, n)| n)
            .sum();
        if n > 0 {
            let _ = writeln!(out, "  |d| in {lo}..{hi}: {n:8}");
        }
    }
    if !r.mismatches.is_empty() {
        let _ = writeln!(out, "\nfirst mismatches:");
        for m in r.mismatches.iter().take(16) {
            let _ = writeln!(
                out,
                "  tfrag {} entry {} channel {}: ours {} ram {}",
                m.tfrag, m.entry, m.channel, m.ours, m.ram
            );
        }
    }
    for b in &r.breakdowns {
        let Some(t) = &b.terms else {
            let _ = writeln!(
                out,
                "\ntfrag {} entry {}: no light set for select {:#06x}; the pass yields transparent\
               \n    black, the game wrote {:?}\n    (record base {:?}, disc placeholder {:?})",
                b.tfrag, b.entry, b.select, b.ram, b.base, b.disc
            );
            continue;
        };
        let _ = writeln!(
            out,
            "\ntfrag {} entry {}: ours {:?} ram {:?}   (select {:#06x}, set {:?}, blended {})",
            b.tfrag, b.entry, t.out, b.ram, t.select, t.set_index, t.blended
        );
        let _ = writeln!(
            out,
            "  normal (az {}, el {}) = {:.3} {:.3} {:.3}",
            t.azimuth, t.elevation, t.normal[0], t.normal[1], t.normal[2]
        );
        let _ = writeln!(
            out,
            "  base {:?}   A colour {:.3} {:.3} {:.3} (w {:.3})  B colour {:.3} {:.3} {:.3} (w {:.3})",
            t.base,
            t.color_a[0], t.color_a[1], t.color_a[2], t.wa,
            t.color_b[0], t.color_b[1], t.color_b[2], t.wb
        );
        let _ = writeln!(
            out,
            "  dA {:.3} -> max {:.3}   dB {:.3} -> max {:.3}",
            t.dot_a, t.max_a, t.dot_b, t.max_b
        );
        let _ = writeln!(
            out,
            "  dir A {:.3} {:.3} {:.3}   dir B {:.3} {:.3} {:.3}",
            t.dir_a[0], t.dir_a[1], t.dir_a[2], t.dir_b[0], t.dir_b[1], t.dir_b[2]
        );
    }
    Ok(())
}

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

    #[test]
    fn the_invariant_mask_excludes_the_pointer_and_the_occlusion_index() {
        let disc = vec![0u8; 0x40];
        let mut ram = disc.clone();
        // The relocated pointer: any difference is allowed.
        ram[0x10..0x14].fill(0xff);
        // The occlusion index, which the loader renumbers.
        ram[0x3a] = 1;
        ram[0x3b] = 2;
        assert!(invariant_equal(&disc, &ram));
        // But the bounding sphere is load-invariant and must match.
        let mut bad = ram.clone();
        bad[1] = 0x99;
        assert!(!invariant_equal(&disc, &bad));
    }

    #[test]
    fn a_ram_image_must_be_exactly_32mib() {
        let image = vec![0u8; crate::ee::EE_SIZE];
        assert!(EeRam::new(&image).is_ok());
        assert!(EeRam::new(&vec![0u8; 1024]).is_err());
    }
}