xenia-rs/crates/xenia-jit/src/lib.rs

//! PPC→x64 block JIT for xenia-rs (`XENIA_JIT`, default OFF).
//!
//! **Phase 0 — skeleton.** This lands the whole runtime substrate (the
//! compiled-block ABI, the interpreter-fallback helper, per-instruction counter
//! bumps, block-exit semantics, and a per-slot code cache) while porting
//! **zero** opcodes to native code: every guest instruction is emitted as a
//! `call` into the interpreter (`xenia_cpu::interpreter::interpret_one`). This
//! makes a JIT-compiled block **byte-identical** to `step_block` by
//! construction, so the golden regression proves the ABI before any opcode is
//! hand-written. Later phases replace individual `call interpret_one` sites
//! with native x64 for the hot opcodes; un-ported opcodes keep falling back.
//!
//! ## Design (context-threading, dynasm-rs)
//! Guest state lives in `PpcContext`; a compiled block is an
//! `extern "C" fn(*mut JitEnv) -> u32` returning a [`StepResult`] discriminant
//! (0 = `Continue`). Emitted code keeps the `PpcContext` pointer in `r15` and
//! the `JitEnv` pointer in `rbx` (both callee-saved, so they survive the helper
//! `call`s). Memory access + un-ported opcodes go through `extern "C"` helpers
//! that receive `JitEnv` and reconstruct `&dyn MemoryAccess` from the fat raw
//! pointer stored in it — no fat-pointer transmute.
//!
//! ## Determinism (the load-bearing invariant)
//! After **every** retired instruction (native or fallback) the block bumps
//! `ctx.cycle_count` and `ctx.timebase` by 1, matching
//! `interpreter.rs::step_block` exactly. Blocks stop at the same instruction
//! the interpreter would (non-`Continue` result, or a taken branch that makes
//! `pc != expected_next`). The JIT code cache mirrors the interpreter block
//! cache's `(start_pc, page_version)` invalidation, and each compiled block
//! **owns a copy** of its decoded instructions so baked instruction pointers
//! can never dangle after a block-cache eviction.

use dynasmrt::{DynasmApi, DynasmLabelApi, dynasm};

use xenia_cpu::block_cache::DecodedBlock;
use xenia_cpu::context::PpcContext;
use xenia_cpu::decoder::DecodedInstr;
use xenia_cpu::interpreter::{StepResult, interpret_one};
use xenia_memory::MemoryAccess;

mod emit;
#[cfg(test)]
mod tests;

/// Runtime environment handed to a compiled block. The emitted prologue reads
/// only `ctx` (via `offset_of!`); `mem` and `last_result` are touched solely by
/// the Rust helpers. `mem` is a real fat raw pointer, so no transmute is needed
/// to reconstruct the `&dyn` in the helpers.
pub struct JitEnv {
    /// Guest CPU state — loaded into `r15` by the block prologue.
    ctx: *mut PpcContext,
    /// The guest memory the block runs against (fat raw pointer). Only the
    /// `extern "C"` helpers dereference this.
    mem: *const dyn MemoryAccess,
    /// The exact `StepResult` of the last instruction the block ran. The Rust
    /// wrapper reads this on a non-`Continue` exit so the full payload (e.g.
    /// `Unimplemented(op)`) is preserved without serializing it through the
    /// `u32` return channel.
    last_result: StepResult,
}

/// A compiled block's callable form. First arg (`rdi`) is the `JitEnv`; the
/// return value (`eax`) is a [`StepResult`] discriminant (0 = `Continue`).
type JitBlockFn = unsafe extern "C" fn(*mut JitEnv) -> u32;

/// Map a `StepResult` to the block's `u32` return channel. Only the
/// `Continue == 0` vs non-zero distinction is load-bearing (the wrapper reads
/// `JitEnv::last_result` for the actual non-`Continue` value); the specific
/// codes are for clarity/debugging.
#[inline]
fn sr_code(r: StepResult) -> u32 {
    match r {
        StepResult::Continue => 0,
        StepResult::SystemCall => 1,
        StepResult::Unimplemented(_) => 2,
        StepResult::Trap => 3,
        StepResult::Halted => 4,
        StepResult::Yield => 5,
    }
}

/// Interpreter fallback for one instruction, called from emitted code.
///
/// SAFETY: invoked only from a compiled block created by [`compile_block`],
/// which passes a `JitEnv` that is live on [`run_jit_block`]'s stack and an
/// `instr` pointing into the owning `CompiledBlock`'s instruction copy (kept
/// alive for the duration of the call). Reconstructs the `&mut PpcContext` and
/// `&dyn MemoryAccess` from the env. Does not bump counters — the block does.
unsafe extern "C" fn jit_interpret_one(env: *mut JitEnv, instr: *const DecodedInstr) -> u32 {
    // SAFETY: see function contract.
    let env = unsafe { &mut *env };
    let ctx = unsafe { &mut *env.ctx };
    let mem: &dyn MemoryAccess = unsafe { &*env.mem };
    let instr = unsafe { &*instr };
    let r = interpret_one(ctx, mem, instr);
    env.last_result = r;
    sr_code(r)
}

// ---- memory-access helpers called from emitted load/store code ----
//
// Each reconstructs `&dyn MemoryAccess` (and, for stores, `&PpcContext`) from
// the `JitEnv` and calls the SAME trait methods the interpreter uses, so MMIO
// routing, mem-watch, and page-version bumps are byte-identical. Loads are pure
// reads; stores replicate the interpreter store arms' reservation-invalidation
// prologue exactly (a no-op when no reservation table is installed).

/// SAFETY (all helpers): `env` is the live `JitEnv` from [`run_jit_block`];
/// `env.ctx`/`env.mem` are valid for the block call. Big-endian handling lives
/// in the trait methods, matching the interpreter.
unsafe extern "C" fn jit_read_u8(env: *mut JitEnv, addr: u32) -> u8 {
    let mem: &dyn MemoryAccess = unsafe { &*(&*env).mem };
    mem.read_u8(addr)
}
unsafe extern "C" fn jit_read_u16(env: *mut JitEnv, addr: u32) -> u16 {
    let mem: &dyn MemoryAccess = unsafe { &*(&*env).mem };
    mem.read_u16(addr)
}
unsafe extern "C" fn jit_read_u32(env: *mut JitEnv, addr: u32) -> u32 {
    let mem: &dyn MemoryAccess = unsafe { &*(&*env).mem };
    mem.read_u32(addr)
}
unsafe extern "C" fn jit_read_u64(env: *mut JitEnv, addr: u32) -> u64 {
    let mem: &dyn MemoryAccess = unsafe { &*(&*env).mem };
    mem.read_u64(addr)
}

/// Reservation invalidation shared by all store arms (mirrors
/// `interpreter.rs`: invalidate a same-line reservation before a write; no-op
/// unless a reservation table is installed and enabled with active reservers).
#[inline]
unsafe fn store_reservation_invalidate(env: &JitEnv, ea: u32) {
    let ctx = unsafe { &*env.ctx };
    if let Some(t) = ctx.reservation_table.as_ref().filter(|t| t.is_enabled()) {
        if t.has_active_reservers() {
            t.invalidate_for_write(ea);
        }
    }
}

unsafe extern "C" fn jit_store_u8(env: *mut JitEnv, addr: u32, val: u64) {
    let env = unsafe { &*env };
    unsafe { store_reservation_invalidate(env, addr) };
    let mem: &dyn MemoryAccess = unsafe { &*env.mem };
    mem.write_u8(addr, val as u8);
}
unsafe extern "C" fn jit_store_u16(env: *mut JitEnv, addr: u32, val: u64) {
    let env = unsafe { &*env };
    unsafe { store_reservation_invalidate(env, addr) };
    let mem: &dyn MemoryAccess = unsafe { &*env.mem };
    mem.write_u16(addr, val as u16);
}
unsafe extern "C" fn jit_store_u32(env: *mut JitEnv, addr: u32, val: u64) {
    let env = unsafe { &*env };
    unsafe { store_reservation_invalidate(env, addr) };
    let mem: &dyn MemoryAccess = unsafe { &*env.mem };
    mem.write_u32(addr, val as u32);
}
unsafe extern "C" fn jit_store_u64(env: *mut JitEnv, addr: u32, val: u64) {
    let env = unsafe { &*env };
    unsafe { store_reservation_invalidate(env, addr) };
    let mem: &dyn MemoryAccess = unsafe { &*env.mem };
    mem.write_u64(addr, val);
}

/// Absolute addresses of the memory helpers, baked into emitted code.
pub(crate) struct MemHelpers {
    pub read_u8: i64,
    pub read_u16: i64,
    pub read_u32: i64,
    pub read_u64: i64,
    pub store_u8: i64,
    pub store_u16: i64,
    pub store_u32: i64,
    pub store_u64: i64,
}
impl MemHelpers {
    pub(crate) fn resolve() -> Self {
        MemHelpers {
            read_u8: jit_read_u8 as usize as i64,
            read_u16: jit_read_u16 as usize as i64,
            read_u32: jit_read_u32 as usize as i64,
            read_u64: jit_read_u64 as usize as i64,
            store_u8: jit_store_u8 as usize as i64,
            store_u16: jit_store_u16 as usize as i64,
            store_u32: jit_store_u32 as usize as i64,
            store_u64: jit_store_u64 as usize as i64,
        }
    }
}

/// One JIT-compiled block. Owns everything the emitted code references so the
/// code, its instruction pointers, and its cache-key metadata share one
/// lifetime.
struct CompiledBlock {
    start_pc: u32,
    /// `DecodedBlock::page_version` at compile time; mismatch on lookup forces
    /// recompilation (mirrors the interpreter block cache invalidation).
    page_version: u64,
    /// Owned copy of the block's decoded instructions. The emitted `call`s bake
    /// raw pointers to these elements, so this boxed slice (stable address)
    /// MUST outlive `func`. Kept alive as a field; not read directly.
    _instrs: Box<[DecodedInstr]>,
    /// Backing executable mapping for `func`. Kept alive as a field.
    _buf: dynasmrt::ExecutableBuffer,
    /// Entry point into `_buf`.
    func: JitBlockFn,
}

// SAFETY: `CompiledBlock` is only ever created, stored, and invoked on the
// single owning HW-slot thread (each `WorkerCtx` has its own `JitCache`), the
// same discipline as the interpreter's per-slot `BlockCache`. The raw pointers
// it holds are self-owned. It is never shared across threads.
unsafe impl Send for CompiledBlock {}

/// Compile `block` into a `CompiledBlock`. Phase 0: every instruction is a
/// `call jit_interpret_one` + the mandatory counter/exit postlude.
fn compile_block(block: &DecodedBlock) -> CompiledBlock {
    // Own the instruction stream first, then bake pointers into the *owned*
    // copy (its addresses are final once boxed).
    let instrs: Box<[DecodedInstr]> = block.instrs.clone().into_boxed_slice();

    // Field offsets + helper addresses resolved at compile time.
    let off = emit::Offsets::resolve();
    let mem_helpers = MemHelpers::resolve();
    let helper = jit_interpret_one as usize as i64;

    let mut ops = dynasmrt::x64::Assembler::new().expect("dynasm assembler");
    let entry = ops.offset();
    let l_exit = ops.new_dynamic_label();
    let l_cont = ops.new_dynamic_label();

    // Prologue: save callee-saved regs we use, keep the stack 16-aligned before
    // the helper calls (entry rsp%16==8; two pushes -> 8; `sub 8` -> 0), pin
    // env in rbx and ctx in r15.
    dynasm!(ops
        ; .arch x64
        ; push rbx
        ; push r15
        ; sub rsp, 8
        ; mov rbx, rdi
        ; mov r15, [rbx + off.env_ctx]
    );

    // Counter/pc deferral: native ops just accumulate `state.pending`; pc and the
    // counters are materialized only at observability points (fallbacks, branch
    // edges, block end). `tail_pc` = Some(next_pc) when the last emitted op was
    // straight-line native (its pc write is deferred to the fall-through end);
    // None when a branch/fallback already set pc.
    let mut state = emit::EmitState::new();
    let mut tail_pc: Option<u32> = None;
    for instr in instrs.iter() {
        match emit::try_emit_native(&mut ops, &off, &mem_helpers, &mut state, instr) {
            // Native non-branch: computation only; pc/counters deferred.
            emit::Emit::Native => {
                tail_pc = Some(instr.addr.wrapping_add(4));
                continue;
            }
            // Native branch: it set pc absolutely and flushed the counters.
            // Append the pc-discontinuity check (taken -> exit via l_cont;
            // fall-through -> continue), matching step_block.
            emit::Emit::Branch => {
                let expected_next = instr.addr.wrapping_add(4) as i32;
                dynasm!(ops
                    ; .arch x64
                    ; cmp DWORD [r15 + off.pc], expected_next
                    ; jne =>l_cont
                );
                tail_pc = None;
                continue;
            }
            // Un-ported opcode: emit the interpreter fallback below.
            emit::Emit::Fallback => {}
        }
        // Interpreter fallback for un-ported opcodes. Make counters + pc current
        // first (the callee may read timebase via mftb and reads/writes pc), then
        // run it and account its own retirement (interpreter order: after execute).
        state.flush_counters(&mut ops, &off);
        let addr = instr.addr as i32;
        let instr_ptr = instr as *const DecodedInstr as usize as i64;
        let expected_next = instr.addr.wrapping_add(4) as i32;
        dynasm!(ops
            ; .arch x64
            ; mov DWORD [r15 + off.pc], addr
            // fallback: eax = jit_interpret_one(env, &instr); env.last_result set
            ; mov rdi, rbx
            ; mov rsi, QWORD instr_ptr
            ; mov rax, QWORD helper
            ; call rax
            // determinism postlude: this instruction retired.
            ; inc QWORD [r15 + off.cycle]
            ; inc QWORD [r15 + off.timebase]
            // non-Continue result -> exit returning the discriminant in eax
            ; test eax, eax
            ; jnz =>l_exit
            // taken-branch (pc discontinuity) -> stop the block, return Continue
            ; cmp DWORD [r15 + off.pc], expected_next
            ; jne =>l_cont
        );
        tail_pc = None;
    }

    // Block end (fall-through): flush the deferred counters and materialize the
    // final pc if the tail was straight-line native. Emitted BEFORE l_cont so a
    // taken branch/fallback (which jumps to l_cont) skips it — its pc/counters
    // are already current.
    state.flush_counters(&mut ops, &off);
    if let Some(p) = tail_pc {
        dynasm!(ops ; .arch x64 ; mov DWORD [r15 + off.pc], p as i32);
    }

    // Natural end / discontinuity exit: Continue (eax=0). Shared epilogue.
    dynasm!(ops
        ; .arch x64
        ; =>l_cont
        ; xor eax, eax
        ; =>l_exit
        ; add rsp, 8
        ; pop r15
        ; pop rbx
        ; ret
    );

    let buf = ops.finalize().expect("dynasm finalize");
    // SAFETY: `entry` is a valid offset into `buf`; the emitted code matches
    // the `JitBlockFn` ABI (System V, first arg rdi, return eax).
    let func: JitBlockFn = unsafe { std::mem::transmute::<*const u8, JitBlockFn>(buf.ptr(entry)) };

    CompiledBlock {
        start_pc: block.start_pc,
        page_version: block.page_version,
        _instrs: instrs,
        _buf: buf,
        func,
    }
}

/// Run a compiled block against `ctx`/`mem`, returning the same `StepResult`
/// the interpreter's `step_block` would. The block bumps `cycle_count`/
/// `timebase` and updates `ctx.pc` in place, exactly like the interpreter.
fn run_jit_block(cb: &CompiledBlock, ctx: &mut PpcContext, mem: &dyn MemoryAccess) -> StepResult {
    // Erase the borrow lifetime so it fits `JitEnv::mem` (a raw
    // `*const dyn MemoryAccess`, i.e. `+ 'static`). This is a lifetime-only
    // transmute — the fat-pointer representation is unchanged — and is sound
    // because `env` does not escape: the block runs synchronously and returns
    // before `mem`'s borrow ends.
    let mem_static: &'static dyn MemoryAccess =
        unsafe { std::mem::transmute::<&dyn MemoryAccess, &'static dyn MemoryAccess>(mem) };
    let mut env = JitEnv {
        ctx: ctx as *mut PpcContext,
        mem: mem_static as *const dyn MemoryAccess,
        last_result: StepResult::Continue,
    };
    // SAFETY: `func` is code emitted by `compile_block` for the JitBlockFn ABI;
    // `env` outlives the call; `cb` (and its owned instrs the code references)
    // is borrowed for the whole call.
    let code = unsafe { (cb.func)(&mut env as *mut JitEnv) };
    if code == 0 {
        StepResult::Continue
    } else {
        env.last_result
    }
}

// Matches the interpreter's `BlockCache` (64K direct-mapped, pc-indexed).
const JIT_CACHE_SIZE: usize = 1 << 16;
const JIT_CACHE_MASK: u32 = (JIT_CACHE_SIZE as u32) - 1;

/// Per-HW-slot JIT code cache. Direct-mapped by guest PC, gated on
/// `(start_pc, page_version)` so it invalidates in lock-step with the
/// interpreter block cache (self-modifying / reloaded code recompiles).
pub struct JitCache {
    slots: Box<[Option<CompiledBlock>]>,
    compiles: u64,
    hits: u64,
}

impl Default for JitCache {
    fn default() -> Self {
        Self::new()
    }
}

impl JitCache {
    pub fn new() -> Self {
        let mut v: Vec<Option<CompiledBlock>> = Vec::with_capacity(JIT_CACHE_SIZE);
        v.resize_with(JIT_CACHE_SIZE, || None);
        Self {
            slots: v.into_boxed_slice(),
            compiles: 0,
            hits: 0,
        }
    }

    pub fn compiles(&self) -> u64 {
        self.compiles
    }
    pub fn hits(&self) -> u64 {
        self.hits
    }

    /// Look up (or compile) the JIT block for `block` and run it. `block` is the
    /// freshly-validated `DecodedBlock` from the interpreter cache, so its
    /// `page_version` is current; we key on it directly.
    pub fn run_or_compile(
        &mut self,
        block: &DecodedBlock,
        ctx: &mut PpcContext,
        mem: &dyn MemoryAccess,
    ) -> StepResult {
        let idx = ((block.start_pc >> 2) & JIT_CACHE_MASK) as usize;
        let fresh = matches!(
            &self.slots[idx],
            Some(cb) if cb.start_pc == block.start_pc && cb.page_version == block.page_version
        );
        if fresh {
            self.hits += 1;
        } else {
            self.compiles += 1;
            self.slots[idx] = Some(compile_block(block));
        }
        let cb = self.slots[idx].as_ref().expect("just populated");
        run_jit_block(cb, ctx, mem)
    }
}

/// Whether the JIT is enabled this run (`XENIA_JIT=1|true|yes`), cached once.
pub fn env_enabled() -> bool {
    use std::sync::OnceLock;
    static ON: OnceLock<bool> = OnceLock::new();
    *ON.get_or_init(|| {
        std::env::var("XENIA_JIT")
            .ok()
            .map(|v| {
                let v = v.trim().to_ascii_lowercase();
                v == "1" || v == "true" || v == "yes"
            })
            .unwrap_or(false)
    })
}