//! MSVC vtable + RTTI detection.
//!
//! Heuristic two-pass scan over the binary's read-only data sections. Pass 1
//! finds candidate vtables — runs of ≥3 contiguous big-endian u32 values that
//! all land on known function entries. Pass 2 attempts the MSVC RTTI walk
//! `vtable[-1] → CompleteObjectLocator → TypeDescriptor → mangled name`. When
//! RTTI is stripped (typical for shipped game binaries), each anonymous vtable
//! gets a deterministic name `ANON_Class_<hex>` keyed by a hash of its
//! sorted method PCs (so identical vtables across multiple class instances
//! collapse to one entry).
//!
//! What this module does NOT do:
//! - Vtables in heap-allocated memory (built at runtime by ctors) are out of
//!   scope — only vtables present statically in `.rdata` / `.data`.
//! - RTTI inheritance (`BaseClassDescriptor` walk) is best-effort; we record
//!   the first-level base list when present and leave it NULL otherwise.
//! - Multiple-inheritance "extra" vftables (one per base subobject) are
//!   detected as independent vtables; we don't link them.
//!
//! Reference: openrce.org "Reversing Microsoft Visual C++" RTTI articles
//! (CompleteObjectLocator / TypeDescriptor / BaseClassDescriptor layout).

use std::collections::BTreeMap;

use xenia_xex::pe::PeSection;

use crate::demangle;

/// Maximum number of consecutive non-function slots tolerated inside an
/// anchor-recovered vtable before the run is considered terminated. MSVC
/// vtables can carry null / pure-virtual / unrecognised-thunk slots in their
/// head or interior; a small budget lets those through without merging two
/// physically-adjacent vtables. Kept small to avoid bridging the gap between
/// distinct tables.
const MAX_ANCHOR_GAP: usize = 2;

/// One detected vtable.
#[derive(Debug, Clone)]
pub struct Vtable {
    /// Absolute VA of `vtable[0]` (first method slot).
    pub address: u32,
    /// Number of methods in the vtable.
    pub length: u32,
    /// Absolute VA of the `CompleteObjectLocator` from `vtable[-1]`, if it
    /// looked like a valid pointer into `.rdata`. NULL when no RTTI / stripped.
    pub col_address: Option<u32>,
    /// Class name. Demangled from RTTI when available, otherwise the synthetic
    /// `ANON_Class_<hex>` form.
    pub class_name: String,
    /// True when the COL → TypeDescriptor walk succeeded.
    pub rtti_present: bool,
    /// First-level base class names from `RTTIClassHierarchyDescriptor`, JSON-encoded.
    /// `None` when not parseable.
    pub base_classes_json: Option<String>,
    /// One entry per slot: function VA in `.text`.
    pub methods: Vec<u32>,
}

/// Run the vtable scan + RTTI walk. `function_starts` is the set of valid
/// `.text` function entry VAs from M1's corrected `functions` table.
#[tracing::instrument(skip_all, fields(image_base = format_args!("{:#010x}", image_base)))]
pub fn analyze(
    pe: &[u8],
    image_base: u32,
    sections: &[PeSection],
    function_starts: &std::collections::BTreeSet<u32>,
) -> Vec<Vtable> {
    analyze_with_anchors(pe, image_base, sections, function_starts, &std::collections::BTreeSet::new())
}

/// Like [`analyze`], but additionally recovers vtables whose base address is
/// known a-priori from a constructor vptr-write store (an "anchor"). The
/// contiguity heuristic in pass 1 fragments any vtable whose head region
/// contains words that don't resolve to recognised function entries (null /
/// pure-virtual / unrecognised thunk slots); those vtables are never emitted
/// and the downstream typed-dispatch resolver can't type objects of that
/// class. An anchor is a *content-independent* vtable signal — the ctor
/// literally installs `vtable_base` into `this+0` via
/// `addis/addi (or lis/ori) → stw rX, 0(rThis)` — so for every anchor not
/// already covered by a pass-1 run we synthesise a vtable starting at that
/// base, reading the fnptr-array run while *tolerating* up to
/// [`MAX_ANCHOR_GAP`] consecutive non-function slots before terminating.
///
/// `anchors` are absolute VAs of vtable bases (from
/// [`scan_vptr_write_constants`]). Existing pass-1 vtables are kept unchanged
/// (no regression): an anchor that already coincides with a detected vtable
/// base is skipped, and an anchor that lands *inside* an existing run is also
/// skipped (it's a sub-object pointer, not a fresh table).
#[tracing::instrument(skip_all, fields(image_base = format_args!("{:#010x}", image_base)))]
pub fn analyze_with_anchors(
    pe: &[u8],
    image_base: u32,
    sections: &[PeSection],
    function_starts: &std::collections::BTreeSet<u32>,
    anchors: &std::collections::BTreeSet<u32>,
) -> Vec<Vtable> {
    let started = std::time::Instant::now();
    // Sections we'll scan for vtable bodies.
    let scan_targets: Vec<&PeSection> = sections
        .iter()
        .filter(|s| matches!(s.name.as_str(), ".rdata" | ".data"))
        .collect();

    // Range table for "is this VA in .rdata or .data?"
    let rdata_ranges: Vec<(u32, u32)> = sections
        .iter()
        .filter(|s| s.name == ".rdata")
        .map(|s| (image_base + s.virtual_address, image_base + s.virtual_address + s.virtual_size))
        .collect();

    let mut candidates: Vec<Vtable> = Vec::new();

    for section in scan_targets {
        let va_start = image_base + section.virtual_address;
        let va_end = va_start + section.virtual_size;
        let raw_start = section.virtual_address as usize;
        let raw_end = (section.virtual_address + section.virtual_size) as usize;
        if raw_end > pe.len() { continue; }
        let bytes = &pe[raw_start..raw_end.min(pe.len())];

        let mut i = 0usize;
        while i + 12 <= bytes.len() {
            // Try to start a run at this 4-aligned offset.
            if !i.is_multiple_of(4) { i += 1; continue; }
            let mut run_len = 0usize;
            let mut methods: Vec<u32> = Vec::new();
            let mut j = i;
            while j + 4 <= bytes.len() {
                let val = u32::from_be_bytes([bytes[j], bytes[j + 1], bytes[j + 2], bytes[j + 3]]);
                if function_starts.contains(&val) {
                    methods.push(val);
                    run_len += 1;
                    j += 4;
                } else {
                    break;
                }
            }
            if run_len >= 3 {
                let address = va_start + (i as u32);
                candidates.push(Vtable {
                    address,
                    length: run_len as u32,
                    col_address: None,
                    class_name: synth_anon_name(&methods),
                    rtti_present: false,
                    base_classes_json: None,
                    methods,
                });
                i += run_len * 4;
            } else {
                i += 4;
            }
        }
        let _ = (va_start, va_end);
    }

    // --- Anchor-driven recovery (vptr-write-anchored vtables) ---
    //
    // Build a coverage interval set from pass-1 runs so we don't re-emit a
    // table for an anchor that already lies within an extracted vtable.
    let mut covered: Vec<(u32, u32)> = candidates
        .iter()
        .map(|v| (v.address, v.address + v.length * 4))
        .collect();
    covered.sort_unstable();

    let is_covered = |addr: u32, covered: &[(u32, u32)]| -> bool {
        covered.iter().any(|&(s, e)| addr >= s && addr < e)
    };

    // Section lookup for "which scan target contains this VA?"
    let scan_targets_va: Vec<(u32, u32, usize, usize)> = sections
        .iter()
        .filter(|s| matches!(s.name.as_str(), ".rdata" | ".data"))
        .map(|s| {
            let va = image_base + s.virtual_address;
            (
                va,
                va + s.virtual_size,
                s.virtual_address as usize,
                (s.virtual_address + s.virtual_size) as usize,
            )
        })
        .collect();

    // Cap a recovered run at the *next anchor* so two physically-adjacent
    // anchored vtables don't merge. We deliberately do NOT cap at pass-1
    // fragments: a fragment is a sub-run the contiguity scan carved out of a
    // larger table, and the anchor legitimately re-absorbs it (subsumed
    // fragments are removed afterwards).
    let anchor_bases: std::collections::BTreeSet<u32> = anchors.iter().copied().collect();

    let mut recovered = 0usize;
    let mut newly: Vec<Vtable> = Vec::new();
    for &anchor in anchors {
        if is_covered(anchor, &covered) { continue; }
        // Locate the containing .rdata/.data section.
        let Some(&(va_lo, va_hi, raw_lo, raw_hi)) =
            scan_targets_va.iter().find(|&&(lo, hi, _, _)| anchor >= lo && anchor < hi)
        else { continue };
        if anchor % 4 != 0 { continue; }
        let raw_hi = raw_hi.min(pe.len());
        // Read the fnptr-array run starting at the anchor. Tolerate small
        // gaps of non-function slots (null / pure-virtual / unrecognised),
        // but require the run to actually contain at least one real function
        // (otherwise it's just data, not a vtable).
        let next_base = anchor_bases.range((anchor + 4)..).next().copied();
        let mut methods: Vec<u32> = Vec::new();
        let mut gap = 0usize;
        let mut real_fns = 0usize;
        let mut off = (anchor - va_lo) as usize + raw_lo;
        let mut va = anchor;
        while off + 4 <= raw_hi && va < va_hi {
            if let Some(nb) = next_base && va >= nb { break; }
            let val = u32::from_be_bytes([pe[off], pe[off + 1], pe[off + 2], pe[off + 3]]);
            if function_starts.contains(&val) {
                methods.push(val);
                real_fns += 1;
                gap = 0;
            } else {
                // A non-function slot. Keep the slot (so downstream slot
                // indexing stays aligned) but count toward the gap budget.
                gap += 1;
                if gap > MAX_ANCHOR_GAP {
                    // Drop the trailing gap slots — they belong past the
                    // table's end.
                    methods.truncate(methods.len().saturating_sub(gap - 1));
                    break;
                }
                methods.push(val);
            }
            off += 4;
            va += 4;
        }
        // Trim any trailing non-function slots (the table ends at its last
        // real method).
        while methods.last().is_some_and(|&m| !function_starts.contains(&m)) {
            methods.pop();
        }
        if real_fns == 0 || methods.is_empty() { continue; }
        let length = methods.len() as u32;
        newly.push(Vtable {
            address: anchor,
            length,
            col_address: None,
            class_name: synth_anon_name(&methods),
            rtti_present: false,
            base_classes_json: None,
            methods,
        });
        recovered += 1;
    }
    if recovered > 0 {
        // Drop pass-1 fragments fully subsumed by a recovered (anchored)
        // vtable — the anchor base is authoritative and the fragment was a
        // contiguity-scan artifact of the same table. Keep fragments that
        // only partially overlap (defensive; shouldn't happen for true
        // sub-runs) so we never lose method coverage.
        let recovered_spans: Vec<(u32, u32)> =
            newly.iter().map(|v| (v.address, v.address + v.length * 4)).collect();
        candidates.retain(|v| {
            !recovered_spans
                .iter()
                .any(|&(s, e)| v.address >= s && v.address + v.length * 4 <= e)
        });
        candidates.extend(newly);
        tracing::info!(recovered, "vtables recovered from vptr-write anchors");
    }
    let _ = &covered;

    // RTTI walk: for each candidate, look at vtable[-1].
    let pe_image_base = image_base;
    for v in &mut candidates {
        if v.address < 4 { continue; }
        let col_off = (v.address - pe_image_base - 4) as usize;
        if col_off + 4 > pe.len() { continue; }
        let col_ptr = u32::from_be_bytes([pe[col_off], pe[col_off + 1], pe[col_off + 2], pe[col_off + 3]]);
        if col_ptr == 0 { continue; }
        if !is_in_ranges(col_ptr, &rdata_ranges) { continue; }

        // Try to extract the TypeDescriptor mangled-name string.
        if let Some((td_ptr, hierarchy_ptr)) = read_col(pe, image_base, col_ptr)
            && let Some(mangled) = read_typedescriptor_name(pe, image_base, td_ptr, &rdata_ranges)
            && let Some(class) = demangle_rtti_typename(&mangled)
        {
            v.col_address = Some(col_ptr);
            v.class_name = class;
            v.rtti_present = true;
            v.base_classes_json = read_class_hierarchy(pe, image_base, hierarchy_ptr, &rdata_ranges);
        }
    }

    let elapsed_ms = started.elapsed().as_millis() as f64;
    let rtti_count = candidates.iter().filter(|v| v.rtti_present).count();
    metrics::histogram!("analysis.phase_ms", "phase" => "vtables").record(elapsed_ms);
    tracing::info!(
        vtables = candidates.len(),
        rtti = rtti_count,
        anon = candidates.len() - rtti_count,
        elapsed_ms,
        "vtable scan complete"
    );
    candidates
}

fn is_in_ranges(addr: u32, ranges: &[(u32, u32)]) -> bool {
    ranges.iter().any(|&(s, e)| addr >= s && addr < e)
}

/// Read 4 big-endian bytes at absolute VA `addr` from the PE image.
fn read_be_u32(pe: &[u8], image_base: u32, addr: u32) -> Option<u32> {
    let off = addr.wrapping_sub(image_base) as usize;
    if off + 4 > pe.len() { return None; }
    Some(u32::from_be_bytes([pe[off], pe[off + 1], pe[off + 2], pe[off + 3]]))
}

/// Parse a `CompleteObjectLocator` at VA `col`. Returns
/// `(type_descriptor_ptr, class_hierarchy_descriptor_ptr)` on success.
///
/// Layout (32-bit MSVC):
/// ```text
///   +0x00 signature (0 for x86 without /GR-, can be 1)
///   +0x04 offset within complete object
///   +0x08 cdOffset (this-pointer adjuster)
///   +0x0C TypeDescriptor *
///   +0x10 RTTIClassHierarchyDescriptor *
/// ```
fn read_col(pe: &[u8], image_base: u32, col: u32) -> Option<(u32, u32)> {
    let td = read_be_u32(pe, image_base, col + 0x0C)?;
    let chd = read_be_u32(pe, image_base, col + 0x10)?;
    if td == 0 { return None; }
    Some((td, chd))
}

/// Read a TypeDescriptor's mangled-name string at VA `td`.
///
/// Layout: `+0x00` vftable ptr, `+0x04` "spare", `+0x08` zero-terminated
/// mangled name (e.g. `.?AVClassName@@`).
fn read_typedescriptor_name(
    pe: &[u8],
    image_base: u32,
    td: u32,
    rdata_ranges: &[(u32, u32)],
) -> Option<String> {
    if !is_in_ranges(td, rdata_ranges) { return None; }
    let name_va = td + 0x08;
    let off = name_va.wrapping_sub(image_base) as usize;
    if off + 1 > pe.len() { return None; }
    // Read up to 256 bytes or until NUL.
    let mut end = off;
    while end < pe.len().min(off + 256) && pe[end] != 0 { end += 1; }
    if end == off { return None; }
    let s = std::str::from_utf8(&pe[off..end]).ok()?;
    // Sanity: MSVC RTTI names always start with `.?A`.
    if !s.starts_with(".?A") { return None; }
    Some(s.to_string())
}

/// Demangle an RTTI type-name string of the form `.?AVClassName@ns@@`.
/// MSVC convention: leading `.` is the marker for an RTTI string; strip it
/// before passing to the demangler.
fn demangle_rtti_typename(rtti_name: &str) -> Option<String> {
    let stripped = rtti_name.strip_prefix('.')?;
    let raw = msvc_demangler::demangle(stripped, msvc_demangler::DemangleFlags::llvm()).ok()?;
    // Output looks like `class xe::apu::AudioSystem` or `struct foo::Bar`.
    let cls = raw
        .strip_prefix("class ")
        .or_else(|| raw.strip_prefix("struct "))
        .or_else(|| raw.strip_prefix("union "))
        .unwrap_or(&raw);
    Some(cls.to_string())
}

/// Best-effort `RTTIClassHierarchyDescriptor` walk: read the
/// `BaseClassArray` entries and demangle each base's TypeDescriptor name.
/// Returns a JSON array string on success.
///
/// Layout:
/// ```text
///   RTTIClassHierarchyDescriptor:
///     +0x00 signature
///     +0x04 attributes
///     +0x08 numBaseClasses
///     +0x0C BaseClassArray *  (-> array of BaseClassDescriptor *)
///   BaseClassDescriptor:
///     +0x00 TypeDescriptor *
///     +0x04 numContainedBases
///     ...
/// ```
fn read_class_hierarchy(
    pe: &[u8],
    image_base: u32,
    chd: u32,
    rdata_ranges: &[(u32, u32)],
) -> Option<String> {
    if !is_in_ranges(chd, rdata_ranges) { return None; }
    let num_bases = read_be_u32(pe, image_base, chd + 0x08)?;
    if num_bases == 0 || num_bases > 256 { return None; } // sanity cap
    let bca_ptr = read_be_u32(pe, image_base, chd + 0x0C)?;
    if !is_in_ranges(bca_ptr, rdata_ranges) { return None; }

    let mut names: Vec<String> = Vec::new();
    for i in 0..num_bases {
        let bcd_ptr = match read_be_u32(pe, image_base, bca_ptr + i * 4) {
            Some(p) if is_in_ranges(p, rdata_ranges) => p,
            _ => return None,
        };
        let td_ptr = match read_be_u32(pe, image_base, bcd_ptr) {
            Some(p) if is_in_ranges(p, rdata_ranges) => p,
            _ => return None,
        };
        let mangled = match read_typedescriptor_name(pe, image_base, td_ptr, rdata_ranges) {
            Some(s) => s,
            None => return None,
        };
        let cls = demangle_rtti_typename(&mangled).unwrap_or(mangled);
        names.push(cls);
    }
    serde_json::to_string(&names).ok()
}

/// Pre-pass: discover candidate vtable *bases* from constructor vptr-write
/// stores, independent of the static contiguity heuristic. A vptr install is
/// the canonical `addis/addi` (or `lis/ori`) immediate build of a constant
/// pointing into `.rdata` / `.data`, followed by `stw rX, 0(rThis)` — i.e. the
/// ctor writing the vtable pointer to `this+0`. We return the set of such
/// constants; these are fed to [`analyze_with_anchors`] so a vtable with
/// non-function head words isn't lost.
///
/// We only consider stores at displacement 0 (the primary vptr; secondary
/// MI vptrs land at non-zero offsets and are handled by the existing
/// contiguity scan / typed-dispatch resolver well enough). The register
/// tracker mirrors the lis+addi propagation used elsewhere and is reset at
/// every basic-block boundary (`block_boundaries`).
pub fn scan_vptr_write_constants(
    pe: &[u8],
    image_base: u32,
    functions: &std::collections::BTreeMap<u32, (u32, bool)>, // start -> (end, is_saverestore)
    sections: &[PeSection],
    block_boundaries: &std::collections::HashSet<u32>,
) -> std::collections::BTreeSet<u32> {
    // Ranges that a vtable base may legitimately live in.
    let data_ranges: Vec<(u32, u32)> = sections
        .iter()
        .filter(|s| matches!(s.name.as_str(), ".rdata" | ".data"))
        .map(|s| (image_base + s.virtual_address, image_base + s.virtual_address + s.virtual_size))
        .collect();
    let in_data = |a: u32| data_ranges.iter().any(|&(s, e)| a >= s && a < e);

    const OP_ADDI: u32 = 14;
    const OP_ADDIS: u32 = 15;
    const OP_ORI: u32 = 24;
    const OP_STW: u32 = 36;
    const OP_X_FORM: u32 = 31;

    let read = |addr: u32| -> Option<u32> {
        let off = addr.wrapping_sub(image_base) as usize;
        if off + 4 > pe.len() { return None; }
        Some(u32::from_be_bytes([pe[off], pe[off + 1], pe[off + 2], pe[off + 3]]))
    };

    let mut anchors: std::collections::BTreeSet<u32> = std::collections::BTreeSet::new();
    for (&fn_start, &(fn_end, is_saverestore)) in functions {
        if is_saverestore { continue; }
        let mut reg: [Option<u32>; 32] = [None; 32];
        let mut pc = fn_start;
        while pc < fn_end {
            if pc != fn_start && block_boundaries.contains(&pc) {
                reg = [None; 32];
            }
            let Some(instr) = read(pc) else { break };
            let op = instr >> 26;
            let rd = ((instr >> 21) & 0x1F) as usize;
            let ra = ((instr >> 16) & 0x1F) as usize;
            let simm = ((instr & 0xFFFF) as i16) as i32;
            let uimm = instr & 0xFFFF;
            match op {
                OP_ADDIS if ra == 0 => reg[rd] = Some(uimm << 16),
                OP_ADDIS => reg[rd] = reg[ra].map(|b| b.wrapping_add(uimm << 16)),
                OP_ADDI if ra != 0 => reg[rd] = reg[ra].map(|b| b.wrapping_add(simm as u32)),
                OP_ADDI => reg[rd] = Some(simm as u32),
                OP_ORI => {
                    let rs = rd;
                    reg[ra] = reg[rs].map(|b| b | uimm);
                }
                OP_STW => {
                    // `stw rS, off(rA)` with displacement 0 = primary vptr install.
                    if ra != 0
                        && simm == 0
                        && let Some(val) = reg[rd]
                        && in_data(val)
                    {
                        anchors.insert(val);
                    }
                }
                32..=35 | 40..=43 | 48..=51 => reg[rd] = None,
                OP_X_FORM => {
                    let xo = (instr >> 1) & 0x3FF;
                    if xo != 444 && xo != 467 { reg[rd] = None; } // keep `or`(444=mr)/`mtspr`-ish
                }
                18 | 16 => {
                    if (instr & 1) != 0 {
                        for r in 0..=12 { reg[r] = None; }
                    }
                }
                _ => {}
            }
            pc = pc.wrapping_add(4);
        }
    }
    anchors
}

/// Synthetic name for an RTTI-stripped vtable, derived from a stable hash of
/// the sorted method-PC list. Two vtables with identical method ordering
/// collapse to the same anonymous name.
fn synth_anon_name(methods: &[u32]) -> String {
    // FNV-1a 64-bit on the sorted PC list; we only use 32 bits for brevity.
    let mut sorted = methods.to_vec();
    sorted.sort_unstable();
    let mut h: u64 = 0xcbf29ce484222325;
    for pc in &sorted {
        for b in pc.to_le_bytes() {
            h ^= b as u64;
            h = h.wrapping_mul(0x100000001b3);
        }
    }
    format!("ANON_Class_{:08X}", (h as u32))
}

/// Build the per-method `(vtable_address, slot, function_address)` list for
/// DB insertion, with optional demangled-name lookup for any function that
/// has a matching `?…` label. Skips slots whose function isn't in the
/// supplied label map.
pub fn methods_table(
    vtables: &[Vtable],
    labels: &std::collections::HashMap<u32, String>,
) -> Vec<(u32, u32, u32, Option<String>, Option<String>)> {
    let mut out = Vec::new();
    for v in vtables {
        for (slot, &fn_va) in v.methods.iter().enumerate() {
            let label = labels.get(&fn_va).cloned();
            let demangled = label.as_ref()
                .and_then(|l| demangle::demangle(l).map(|d| d.raw_demangled));
            out.push((v.address, slot as u32, fn_va, label, demangled));
        }
    }
    out
}

/// Build a `class_name → Vtable` summary for the `classes` table. Multiple
/// vtables sharing the same class name (multiple instances at link time)
/// collapse via `BTreeMap` — the first detected vtable wins.
pub fn classes_table(vtables: &[Vtable]) -> Vec<(String, u32, bool, Option<String>)> {
    let mut by_name: BTreeMap<String, &Vtable> = BTreeMap::new();
    for v in vtables {
        by_name.entry(v.class_name.clone()).or_insert(v);
    }
    by_name
        .into_iter()
        .map(|(name, v)| (name, v.address, v.rtti_present, v.base_classes_json.clone()))
        .collect()
}

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

    #[test]
    fn synth_anon_name_is_stable() {
        let a = synth_anon_name(&[0x82001000, 0x82001100, 0x82001200]);
        let b = synth_anon_name(&[0x82001200, 0x82001000, 0x82001100]);
        assert_eq!(a, b, "anon name must be order-independent");
    }

    #[test]
    fn synth_anon_name_differs_for_different_methods() {
        let a = synth_anon_name(&[0x82001000, 0x82001100]);
        let b = synth_anon_name(&[0x82002000, 0x82002100]);
        assert_ne!(a, b);
    }

    #[test]
    fn detects_3_method_vtable_in_rdata() {
        let image_base = 0x82000000u32;
        let rdata_va = 0x1000u32;
        let text_va = 0x2000u32;
        let rdata_size = 16u32;
        let text_size = 0x100u32;

        // PE buffer big enough for both sections.
        let total = (text_va + text_size) as usize;
        let mut pe = vec![0u8; total];

        // Vtable: 3 method PCs at .rdata start, all valid function entries.
        let m: [u32; 3] = [image_base + text_va, image_base + text_va + 0x10, image_base + text_va + 0x20];
        for (i, val) in m.iter().enumerate() {
            pe[rdata_va as usize + i * 4..rdata_va as usize + (i + 1) * 4]
                .copy_from_slice(&val.to_be_bytes());
        }

        let sections = vec![
            PeSection {
                name: ".rdata".into(),
                virtual_address: rdata_va,
                virtual_size: rdata_size,
                raw_offset: rdata_va,
                raw_size: rdata_size,
                flags: 0x4000_0040,
            },
            PeSection {
                name: ".text".into(),
                virtual_address: text_va,
                virtual_size: text_size,
                raw_offset: text_va,
                raw_size: text_size,
                flags: 0x6000_0020,
            },
        ];
        let mut function_starts = std::collections::BTreeSet::new();
        for &pc in &m { function_starts.insert(pc); }

        let vtables = analyze(&pe, image_base, &sections, &function_starts);
        assert_eq!(vtables.len(), 1);
        assert_eq!(vtables[0].length, 3);
        assert_eq!(vtables[0].address, image_base + rdata_va);
        assert!(vtables[0].class_name.starts_with("ANON_Class_"));
        assert!(!vtables[0].rtti_present);
    }

    #[test]
    fn anchor_recovers_vtable_with_nonfn_head() {
        // A vtable whose head has a null + an unrecognised word, so the
        // contiguity scan (≥3 contiguous known fns) fragments it. The anchor
        // (from a ctor vptr-write) must recover the whole table from its base.
        let image_base = 0x82000000u32;
        let rdata_va = 0x1000u32;
        let text_va = 0x2000u32;
        let rdata_size = 0x40u32;
        let text_size = 0x100u32;
        let total = (text_va + text_size) as usize;
        let mut pe = vec![0u8; total];

        let f0 = image_base + text_va;
        let f1 = image_base + text_va + 0x10;
        let f2 = image_base + text_va + 0x20;
        // Slots: [null, NONFN(0xDEAD), f0, f1, f2]
        let slots: [u32; 5] = [0, 0xDEADBEEF, f0, f1, f2];
        for (i, val) in slots.iter().enumerate() {
            pe[rdata_va as usize + i * 4..rdata_va as usize + (i + 1) * 4]
                .copy_from_slice(&val.to_be_bytes());
        }

        let sections = vec![
            PeSection {
                name: ".rdata".into(),
                virtual_address: rdata_va,
                virtual_size: rdata_size,
                raw_offset: rdata_va,
                raw_size: rdata_size,
                flags: 0x4000_0040,
            },
            PeSection {
                name: ".text".into(),
                virtual_address: text_va,
                virtual_size: text_size,
                raw_offset: text_va,
                raw_size: text_size,
                flags: 0x6000_0020,
            },
        ];
        let mut function_starts = std::collections::BTreeSet::new();
        for &pc in &[f0, f1, f2] { function_starts.insert(pc); }

        // Without an anchor: the head gap (null + nonfn = 2 slots) means the
        // contiguous run is only [f0,f1,f2]=3 starting at +0x08, so pass-1
        // still finds it but at the WRONG base (0x...1008), not the true base.
        let no_anchor = analyze(&pe, image_base, &sections, &function_starts);
        assert!(
            !no_anchor.iter().any(|v| v.address == image_base + rdata_va),
            "without anchor the table is not recovered at its true base"
        );

        // With the anchor at the true base:
        let mut anchors = std::collections::BTreeSet::new();
        anchors.insert(image_base + rdata_va);
        let with_anchor =
            analyze_with_anchors(&pe, image_base, &sections, &function_starts, &anchors);
        let v = with_anchor
            .iter()
            .find(|v| v.address == image_base + rdata_va)
            .expect("anchor must recover vtable at its true base");
        // length spans through f2 (slot 4): 5 slots.
        assert_eq!(v.length, 5, "table spans null/nonfn head through last fn");
        assert_eq!(v.methods[2], f0);
        assert_eq!(v.methods[4], f2);
    }

    #[test]
    fn scan_vptr_write_constants_finds_ctor_store() {
        // Encode a ctor: addis r11,r0,0x8201; addi r11,r11,lo; stw r11,0(r31)
        // installing vtable base 0x8200A908 into this+0.
        let image_base = 0x82000000u32;
        let ctor = 0x82001000u32;
        let mut pe = vec![0u8; 0x4000];
        // Lay out a tiny .rdata at 0x...A900 so the constant lands in-range.
        let vt_base = 0x8200A908u32; // 0x82010000 - 22264
        let addis = (15u32 << 26) | (11 << 21) | (0 << 16) | 0x8201;
        let lo = (vt_base & 0xFFFF) as i16; // -22264
        let addi = (14u32 << 26) | (11 << 21) | (0 << 16) | ((lo as u16) as u32);
        // addi r11,r0,lo would set r11=lo (sign-extended); we need addis+addi
        // chained. Re-encode addis into r11 from r0, then addi r11,r11,lo.
        let addi2 = (14u32 << 26) | (11 << 21) | (11 << 16) | ((lo as u16) as u32);
        let stw = (36u32 << 26) | (11 << 21) | (31 << 16) | 0; // stw r11,0(r31)
        let at = (ctor - image_base) as usize;
        pe[at..at + 4].copy_from_slice(&addis.to_be_bytes());
        pe[at + 4..at + 8].copy_from_slice(&addi2.to_be_bytes());
        pe[at + 8..at + 12].copy_from_slice(&stw.to_be_bytes());
        let _ = addi;

        let sections = vec![PeSection {
            name: ".rdata".into(),
            virtual_address: 0xA900,
            virtual_size: 0x200,
            raw_offset: 0xA900,
            raw_size: 0x200,
            flags: 0x4000_0040,
        }];
        let mut funcs: std::collections::BTreeMap<u32, (u32, bool)> = std::collections::BTreeMap::new();
        funcs.insert(ctor, (ctor + 0x40, false));
        let anchors = scan_vptr_write_constants(
            &pe, image_base, &funcs, &sections, &std::collections::HashSet::new(),
        );
        assert!(anchors.contains(&vt_base), "ctor vptr store must yield anchor {vt_base:#x}, got {anchors:?}");
    }

    #[test]
    fn rejects_2_method_run() {
        let image_base = 0x82000000u32;
        let rdata_va = 0x1000u32;
        let text_va = 0x2000u32;

        let total = (text_va + 0x100) as usize;
        let mut pe = vec![0u8; total];
        let m: [u32; 2] = [image_base + text_va, image_base + text_va + 0x10];
        for (i, val) in m.iter().enumerate() {
            pe[rdata_va as usize + i * 4..rdata_va as usize + (i + 1) * 4]
                .copy_from_slice(&val.to_be_bytes());
        }
        let sections = vec![
            PeSection {
                name: ".rdata".into(),
                virtual_address: rdata_va,
                virtual_size: 8,
                raw_offset: rdata_va,
                raw_size: 8,
                flags: 0x4000_0040,
            },
            PeSection {
                name: ".text".into(),
                virtual_address: text_va,
                virtual_size: 0x100,
                raw_offset: text_va,
                raw_size: 0x100,
                flags: 0x6000_0020,
            },
        ];
        let mut function_starts = std::collections::BTreeSet::new();
        for &pc in &m { function_starts.insert(pc); }
        let vtables = analyze(&pe, image_base, &sections, &function_starts);
        assert_eq!(vtables.len(), 0, "runs of 2 must be rejected to keep false-positive rate down");
    }
}