xenia-rs/crates/xenia-cpu/src/recompiler.rs

//! PPC block recompiler (JIT) — **M0 foundation**.
//!
//! This is the seam a staged block-recompiler plugs into. M0 ships **no real
//! codegen**: [`run_block`] executes a block by falling back to the
//! interpreter's [`crate::interpreter::execute`] for every instruction, so it
//! is bit-identical to [`crate::interpreter::step_block`] by construction. Its
//! purpose is to prove the integration seam and the **differential harness**
//! before any lowering exists. Later stages replace the per-instruction
//! fallback with lowered ops (closure-threaded, then Cranelift machine code)
//! opcode-by-opcode, each gated through the harness below.
//!
//! ## Determinism is the whole game
//! `cycle_count`/`timebase` are bumped once per guest instruction and drive the
//! IPM clock the movie/scheduler depend on. [`run_block`] reproduces
//! `step_block`'s counting and stop semantics **exactly**. The interpreter is
//! the permanent reference oracle; the recompiler must never diverge from it.
//!
//! ## Differential harness (`XENIA_JIT_DIFF`) — in-process
//! The guest is only *coarsely* deterministic: `coord_idle_advance` ticks vsync
//! from wall-clock when the scheduler idles, so vblank ISRs land at
//! wall-clock-dependent instruction boundaries and two separate runs are **not**
//! bit-identical. Comparing across runs would measure that jitter, not JIT
//! divergence. So the harness is **in-process** ([`diff_step`]):
//!
//!   * the **interpreter is authoritative** — it drives the real context and
//!     memory, so a JIT bug can never corrupt the run;
//!   * the **JIT runs speculatively** on a *clone* of the pre-block context
//!     against an [`OverlayMemory`] (writes buffered, reads fall through to real
//!     *pre-block* memory), then its registers are compared to the
//!     interpreter's;
//!   * blocks that touch **MMIO** (a device callback can't be run twice) or are
//!     **`sync_sensitive`** (reservation/barrier ops share cross-thread state)
//!     are skipped, not compared.
//!
//! Run with `XENIA_JIT_DIFF=1`; [`report_diff_summary`] prints checked/skipped/
//! mismatch counts at exit. Zero mismatches ⇒ the recompiler matches the
//! interpreter on every comparable block.

use std::cell::{Cell, RefCell};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};

use crate::block_cache::DecodedBlock;
use crate::context::PpcContext;
use crate::interpreter::{execute, step_block, StepResult};
use crate::jit::{JitCache, MemEnv, RET_CONTINUE};
use crate::opcode::PpcOpcode;
use xenia_memory::MemoryAccess;

/// Cached env gate. 0 = uninitialised, 1 = on, 2 = off.
#[inline]
fn cached_flag(cell: &AtomicU8, var: &str) -> bool {
    match cell.load(Ordering::Relaxed) {
        1 => true,
        2 => false,
        _ => {
            let on = std::env::var_os(var).is_some();
            cell.store(if on { 1 } else { 2 }, Ordering::Relaxed);
            on
        }
    }
}

/// `XENIA_JIT` — route block execution through the recompiler instead of the
/// interpreter's `step_block`. Opt-in; cheap cached check on the hot path.
#[inline]
pub fn jit_enabled() -> bool {
    static F: AtomicU8 = AtomicU8::new(0);
    cached_flag(&F, "XENIA_JIT")
}

/// `XENIA_JIT_DIFF` — enable the in-process differential harness ([`diff_step`]).
#[inline]
pub fn diff_enabled() -> bool {
    static F: AtomicU8 = AtomicU8::new(0);
    cached_flag(&F, "XENIA_JIT_DIFF")
}

/// `XENIA_JIT_CHAIN` — enable native superblock chaining (the JIT compiles a
/// chain of same-page blocks into one function with direct block-to-block
/// jumps). This is only the *env request*; the scheduler additionally requires
/// the diff harness OFF and no diagnostic probes / mem-watch armed before it
/// actually enables chaining (see the config install in `run_superblock`).
#[inline]
pub fn chain_requested() -> bool {
    static F: AtomicU8 = AtomicU8::new(0);
    cached_flag(&F, "XENIA_JIT_CHAIN")
}

// ---- opcode histogram (XENIA_JIT_HIST) — data-drives M1 coverage ---------

#[inline]
fn hist_enabled() -> bool {
    static F: AtomicU8 = AtomicU8::new(0);
    cached_flag(&F, "XENIA_JIT_HIST")
}

thread_local! {
    static HIST: RefCell<HashMap<PpcOpcode, u64>> = RefCell::new(HashMap::new());
    static HIST_TICK: Cell<u32> = const { Cell::new(0) };
}

/// Sample 1-in-16 instructions — relative opcode frequencies are unbiased at
/// this rate and the HashMap cost stops dominating the run.
#[inline]
fn hist_tally(op: PpcOpcode) {
    let t = HIST_TICK.with(|c| {
        let n = c.get().wrapping_add(1);
        c.set(n);
        n
    });
    if t & 0xF == 0 {
        HIST.with(|h| *h.borrow_mut().entry(op).or_insert(0) += 1);
    }
}

/// Print the opcode-frequency histogram (top 40 + cumulative %), so M1 coverage
/// targets the opcodes that actually dominate boot+movie execution.
pub fn report_histogram() {
    if !hist_enabled() {
        return;
    }
    let mut v: Vec<(PpcOpcode, u64)> =
        HIST.with(|h| h.borrow().iter().map(|(&k, &c)| (k, c)).collect());
    v.sort_by(|a, b| b.1.cmp(&a.1));
    let total: u64 = v.iter().map(|(_, c)| c).sum();
    eprintln!("=== JIT opcode histogram (total={total}, {} distinct) ===", v.len());
    let mut cum = 0u64;
    for (i, (op, c)) in v.iter().take(40).enumerate() {
        cum += c;
        eprintln!(
            "  {:>2}. {:<14?} {:>12} {:>5.1}%  cum {:>5.1}%",
            i + 1,
            op,
            c,
            100.0 * *c as f64 / total as f64,
            100.0 * cum as f64 / total as f64
        );
    }
}

// Coverage telemetry: blocks executed as native code vs interpreted (uncovered).
static JIT_COMPILED_RUN: AtomicU64 = AtomicU64::new(0);
static JIT_INTERP_RUN: AtomicU64 = AtomicU64::new(0);

/// Execute one decoded block, natively if it is fully covered, else via the
/// interpreter fallback.
///
/// If [`JitCache::get_or_compile`] returns a compiled entry point, the whole
/// block runs as native code (which advances `pc`/`cycle_count`/`timebase`
/// exactly like the interpreter — validated by the diff harness). Otherwise
/// this is byte-for-byte the same loop as [`crate::interpreter::step_block`]:
/// bump `cycle_count`/`timebase` per instruction, bail on the first
/// non-`Continue` result, and stop on a PC discontinuity (only the terminator
/// may branch).
pub fn run_block(
    ctx: &mut PpcContext,
    mem: &dyn MemoryAccess,
    block: &DecodedBlock,
    jit: &mut JitCache,
    remaining_budget: u64,
) -> StepResult {
    if let Some(f) = jit.get_or_compile(block, mem) {
        // NB: under superblock chaining one call runs *several* guest blocks,
        // so this counter (and the coverage %) undercounts native guest blocks
        // vs the per-block `JIT_INTERP_RUN`. It stays a native-vs-interpreted
        // *call* ratio; treat the chained coverage % as a floor.
        JIT_COMPILED_RUN.fetch_add(1, Ordering::Relaxed);
        let env = jit.mem_env(mem);
        // `remaining_budget` bounds a superblock's internal chaining; a
        // single-block compilation ignores it. A compiled (super)block always
        // runs to completion and returns `Continue` — covered ops never branch
        // out / fault, and the chain only ever stops at a clean block boundary.
        let raw = f(ctx as *mut PpcContext, &env as *const MemEnv, remaining_budget);
        debug_assert_eq!(raw, RET_CONTINUE, "covered block returned non-Continue");
        return StepResult::Continue;
    }
    JIT_INTERP_RUN.fetch_add(1, Ordering::Relaxed);

    let hist = hist_enabled();
    let mut result = StepResult::Continue;
    for instr in &block.instrs {
        let expected_next = instr.addr.wrapping_add(4);
        if hist {
            hist_tally(instr.opcode);
        }
        // M0 fallback: identical to the interpreter. Future stages dispatch
        // lowered ops here and only fall back for uncompiled opcodes.
        result = execute(ctx, mem, instr);
        ctx.cycle_count += 1;
        ctx.timebase += 1;
        if !matches!(result, StepResult::Continue) {
            return result;
        }
        if ctx.pc != expected_next {
            break;
        }
    }
    result
}

// ---- in-process differential harness -------------------------------------

static CHECKED: AtomicU64 = AtomicU64::new(0);
static SKIPPED: AtomicU64 = AtomicU64::new(0);
static MISMATCH: AtomicU64 = AtomicU64::new(0);
const MAX_MISMATCH_PRINTS: u64 = 20;

/// Run one block under the differential harness. The **interpreter is
/// authoritative** (drives `ctx`+`mem`); the JIT runs speculatively on a clone
/// against an overlay and its registers are compared. Returns the interpreter's
/// `StepResult`.
pub fn diff_step(
    ctx: &mut PpcContext,
    mem: &dyn MemoryAccess,
    block: &DecodedBlock,
    jit: &mut JitCache,
) -> StepResult {
    // Reservation/barrier blocks share cross-thread state; don't speculate.
    if block.sync_sensitive {
        SKIPPED.fetch_add(1, Ordering::Relaxed);
        return step_block(ctx, mem, block);
    }

    // Speculative JIT run on a clone against buffered/overlay memory. Runs
    // BEFORE the authoritative interpreter so its reads see pre-block memory.
    let mut cand = ctx.clone();
    cand.reservation_table = None; // never touch the shared reservation table
    let overlay = OverlayMemory::new(mem);
    // Chaining is disabled under the diff harness (a multi-block superblock
    // can't be compared against one interpreter `step_block`), so this always
    // runs a single block; the budget is irrelevant.
    let _ = run_block(&mut cand, &overlay, block, jit, u64::MAX);
    let touched_mmio = overlay.touched_mmio.get();

    // Authoritative interpreter run: commits real ctx + memory.
    let auth = step_block(ctx, mem, block);

    if touched_mmio {
        SKIPPED.fetch_add(1, Ordering::Relaxed);
    } else {
        CHECKED.fetch_add(1, Ordering::Relaxed);
        compare(&cand, ctx, block.start_pc);
    }
    auth
}

/// Compare speculative-JIT (`cand`) vs authoritative-interpreter (`auth`)
/// register state. Reports the first [`MAX_MISMATCH_PRINTS`] divergent blocks.
/// Covers the integer + FP + control state (M0/M1 scope); vector regs join when
/// VMX opcodes are lowered.
fn compare(cand: &PpcContext, auth: &PpcContext, start_pc: u32) {
    let mut diffs: Vec<String> = Vec::new();
    for i in 0..32 {
        if cand.gpr[i] != auth.gpr[i] {
            diffs.push(format!("r{i} jit={:#018x} int={:#018x}", cand.gpr[i], auth.gpr[i]));
        }
    }
    for i in 0..32 {
        if cand.fpr[i].to_bits() != auth.fpr[i].to_bits() {
            diffs.push(format!(
                "f{i} jit={:#018x} int={:#018x}",
                cand.fpr[i].to_bits(),
                auth.fpr[i].to_bits()
            ));
        }
    }
    if cand.lr != auth.lr {
        diffs.push(format!("lr jit={:#x} int={:#x}", cand.lr, auth.lr));
    }
    if cand.ctr != auth.ctr {
        diffs.push(format!("ctr jit={:#x} int={:#x}", cand.ctr, auth.ctr));
    }
    if cand.pc != auth.pc {
        diffs.push(format!("pc jit={:#x} int={:#x}", cand.pc, auth.pc));
    }
    for i in 0..8 {
        if cand.cr[i].as_u8() != auth.cr[i].as_u8() {
            diffs.push(format!("cr{i} jit={:#x} int={:#x}", cand.cr[i].as_u8(), auth.cr[i].as_u8()));
        }
    }
    if (cand.xer_ca, cand.xer_ov, cand.xer_so, cand.xer_tbc)
        != (auth.xer_ca, auth.xer_ov, auth.xer_so, auth.xer_tbc)
    {
        diffs.push("xer".into());
    }
    if cand.fpscr != auth.fpscr {
        diffs.push(format!("fpscr jit={:#x} int={:#x}", cand.fpscr, auth.fpscr));
    }
    if cand.cycle_count != auth.cycle_count {
        diffs.push(format!("cycles jit={} int={}", cand.cycle_count, auth.cycle_count));
    }

    if !diffs.is_empty() {
        let n = MISMATCH.fetch_add(1, Ordering::Relaxed);
        if n < MAX_MISMATCH_PRINTS {
            eprintln!("JIT-DIFF MISMATCH block={start_pc:#010x}: {}", diffs.join(", "));
        }
    }
}

/// Print checked/skipped/mismatch tallies (call at clean exit).
pub fn report_diff_summary() {
    if !diff_enabled() {
        return;
    }
    let (c, s, m) = (
        CHECKED.load(Ordering::Relaxed),
        SKIPPED.load(Ordering::Relaxed),
        MISMATCH.load(Ordering::Relaxed),
    );
    eprintln!(
        "=== JIT-DIFF SUMMARY: checked={c} skipped(mmio/sync)={s} MISMATCHES={m} — {} ===",
        if m == 0 { "CLEAN ✓" } else { "DIVERGENCE ✗" }
    );
}

/// Print native-vs-interpreted block coverage (call at clean exit). Shown
/// whenever the JIT is engaged (`XENIA_JIT` or `XENIA_JIT_DIFF`) so we can see
/// how much of the real boot+movie workload the current covered set captures.
pub fn report_jit_summary() {
    if !jit_enabled() && !diff_enabled() {
        return;
    }
    let compiled = JIT_COMPILED_RUN.load(Ordering::Relaxed);
    let interp = JIT_INTERP_RUN.load(Ordering::Relaxed);
    let total = compiled + interp;
    let pct = if total == 0 { 0.0 } else { 100.0 * compiled as f64 / total as f64 };
    eprintln!(
        "=== JIT COVERAGE: native blocks={compiled} interpreted={interp} ({pct:.2}% of {total} block-runs native) ==="
    );
}

// ---- overlay memory (speculative-write buffer) ---------------------------

/// A [`MemoryAccess`] wrapper for speculative JIT execution: reads fall through
/// to the real (pre-block) memory, writes are buffered in a byte-granular
/// overlay (never committed), and any MMIO touch sets `touched_mmio` so the
/// caller can skip comparing that block. Big-endian byte order throughout,
/// matching `GuestMemory`.
struct OverlayMemory<'a> {
    real: &'a dyn MemoryAccess,
    overlay: RefCell<HashMap<u32, u8>>,
    touched_mmio: Cell<bool>,
}

impl<'a> OverlayMemory<'a> {
    fn new(real: &'a dyn MemoryAccess) -> Self {
        Self { real, overlay: RefCell::new(HashMap::new()), touched_mmio: Cell::new(false) }
    }
    #[inline]
    fn rb(&self, addr: u32) -> u8 {
        match self.overlay.borrow().get(&addr) {
            Some(&b) => b,
            None => self.real.read_u8(addr),
        }
    }
    #[inline]
    fn wb(&self, addr: u32, b: u8) {
        self.overlay.borrow_mut().insert(addr, b);
    }
    /// Returns true (and flags the block) if `addr` is MMIO — callers must not
    /// touch it speculatively.
    #[inline]
    fn is_mmio_stop(&self, addr: u32) -> bool {
        if self.real.is_mmio(addr) {
            self.touched_mmio.set(true);
            true
        } else {
            false
        }
    }
}

impl<'a> MemoryAccess for OverlayMemory<'a> {
    fn read_u8(&self, a: u32) -> u8 {
        if self.is_mmio_stop(a) {
            return 0;
        }
        self.rb(a)
    }
    fn read_u16(&self, a: u32) -> u16 {
        if self.is_mmio_stop(a) {
            return 0;
        }
        ((self.rb(a) as u16) << 8) | self.rb(a.wrapping_add(1)) as u16
    }
    fn read_u32(&self, a: u32) -> u32 {
        if self.is_mmio_stop(a) {
            return 0;
        }
        ((self.rb(a) as u32) << 24)
            | ((self.rb(a.wrapping_add(1)) as u32) << 16)
            | ((self.rb(a.wrapping_add(2)) as u32) << 8)
            | self.rb(a.wrapping_add(3)) as u32
    }
    fn read_u64(&self, a: u32) -> u64 {
        if self.is_mmio_stop(a) {
            return 0;
        }
        let hi = self.read_u32(a) as u64;
        let lo = self.read_u32(a.wrapping_add(4)) as u64;
        (hi << 32) | lo
    }
    fn write_u8(&self, a: u32, v: u8) {
        if self.is_mmio_stop(a) {
            return;
        }
        self.wb(a, v);
    }
    fn write_u16(&self, a: u32, v: u16) {
        if self.is_mmio_stop(a) {
            return;
        }
        self.wb(a, (v >> 8) as u8);
        self.wb(a.wrapping_add(1), v as u8);
    }
    fn write_u32(&self, a: u32, v: u32) {
        if self.is_mmio_stop(a) {
            return;
        }
        self.wb(a, (v >> 24) as u8);
        self.wb(a.wrapping_add(1), (v >> 16) as u8);
        self.wb(a.wrapping_add(2), (v >> 8) as u8);
        self.wb(a.wrapping_add(3), v as u8);
    }
    fn write_u64(&self, a: u32, v: u64) {
        self.write_u32(a, (v >> 32) as u32);
        self.write_u32(a.wrapping_add(4), v as u32);
    }
    // execute() never calls translate/translate_mut for real memory (only the
    // test mock does), so returning None is safe here.
    fn translate(&self, _a: u32) -> Option<*const u8> {
        None
    }
    fn translate_mut(&self, _a: u32) -> Option<*mut u8> {
        None
    }
    fn page_version(&self, a: u32) -> u64 {
        self.real.page_version(a)
    }
    fn is_mmio(&self, a: u32) -> bool {
        self.real.is_mmio(a)
    }
}