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

//! Cranelift machine-code block JIT (Stage-2).
//!
//! Compiles a fully-covered [`DecodedBlock`] to a native function and runs it in
//! place of the interpreter. A block is compiled **only if every opcode in it is
//! covered** (see [`covered`]); otherwise the caller interprets the whole block.
//! Coverage grows opcode-by-opcode, each one validated against the interpreter by
//! the recompiler's differential harness (`XENIA_JIT_DIFF`).
//!
//! ## ABI
//! Compiled block: `extern "C" fn(ctx: *mut PpcContext, mem: *const MemEnv) -> u32`.
//! The return is a [`StepResult`] discriminant ([`RET_CONTINUE`] etc.). Guest
//! registers are loaded/stored directly from the `#[repr(C)]` `PpcContext` at
//! `offset_of!` offsets; memory goes through trampolines that call the live
//! `MemoryAccess` (preserving MMIO / mem-watch / page-version) — added as
//! load/store opcodes are covered.
//!
//! ## Determinism
//! A covered block is straight-line and always runs to completion (covered ops
//! never fault or yield), so `cycle_count`/`timebase` advance by exactly the
//! block's instruction count — matching the interpreter's per-instruction bump.

use core::mem::offset_of;

use cranelift_codegen::ir::{types, AbiParam, InstBuilder, MemFlags, Value};
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{default_libcall_names, Module};

use crate::block_cache::DecodedBlock;
use crate::context::PpcContext;
use crate::decoder::DecodedInstr;
use crate::opcode::PpcOpcode;
use xenia_memory::MemoryAccess;

/// `StepResult` discriminants returned by compiled blocks. Kept in sync with
/// [`crate::interpreter::StepResult`] (only `Continue` is produced by covered
/// blocks today; the rest exist for when branch/sc/trap ops are covered).
pub const RET_CONTINUE: u32 = 0;

/// Thin environment handed to memory trampolines: wraps the live `MemoryAccess`.
#[repr(C)]
pub struct MemEnv<'a> {
    pub mem: &'a dyn MemoryAccess,
}

/// A compiled block: a native entry point valid for as long as the owning
/// [`Jit`]'s module lives (the module owns the executable memory).
pub type CompiledFn = extern "C" fn(*mut PpcContext, *const MemEnv) -> u32;

/// Is this opcode currently lowerable to native code? Must mirror [`emit_op`].
/// A block is JIT-compiled only if `covered` is true for all its instructions.
pub fn covered(instr: &DecodedInstr) -> bool {
    matches!(instr.opcode, PpcOpcode::addi | PpcOpcode::addis)
}

/// True if the whole block can be compiled (all opcodes covered).
pub fn block_covered(block: &DecodedBlock) -> bool {
    block.instrs.iter().all(covered)
}

/// The Cranelift JIT. Owns the module (and thus all compiled code memory), so it
/// must outlive every [`CompiledFn`] it produces.
pub struct Jit {
    module: JITModule,
    ctx: cranelift_codegen::Context,
    fbctx: FunctionBuilderContext,
}

impl Jit {
    pub fn new() -> Result<Self, String> {
        // Native host ISA + default flags. `JITBuilder::new` resolves the host
        // target via cranelift-native (a transitive dep of cranelift-jit).
        let builder = JITBuilder::new(default_libcall_names()).map_err(|e| e.to_string())?;
        let module = JITModule::new(builder);
        let ctx = module.make_context();
        Ok(Self { module, ctx, fbctx: FunctionBuilderContext::new() })
    }

    /// Compile `block` to native code, or return `None` if any opcode is
    /// uncovered (caller interprets the block).
    pub fn compile(&mut self, block: &DecodedBlock) -> Option<CompiledFn> {
        if !block_covered(block) {
            return None;
        }
        let ptr = self.module.target_config().pointer_type();
        self.module.clear_context(&mut self.ctx);
        self.ctx.func.signature.params.push(AbiParam::new(ptr)); // ctx
        self.ctx.func.signature.params.push(AbiParam::new(ptr)); // mem env
        self.ctx.func.signature.returns.push(AbiParam::new(types::I32));

        {
            let mut b = FunctionBuilder::new(&mut self.ctx.func, &mut self.fbctx);
            let entry = b.create_block();
            b.append_block_params_for_function_params(entry);
            b.switch_to_block(entry);
            b.seal_block(entry);
            let ctxp = b.block_params(entry)[0];
            let _memenv = b.block_params(entry)[1];

            for instr in &block.instrs {
                emit_op(&mut b, ctxp, instr);
            }

            // Straight-line covered block: PC advances linearly to end_pc.
            let end_pc = b.ins().iconst(types::I32, block.end_pc as i64);
            b.ins().store(MemFlags::trusted(), end_pc, ctxp, off(offset_of!(PpcContext, pc)));

            // cycle_count += N ; timebase += N (one per executed instruction).
            let n = block.instrs.len() as i64;
            bump_u64(&mut b, ctxp, offset_of!(PpcContext, cycle_count), n);
            bump_u64(&mut b, ctxp, offset_of!(PpcContext, timebase), n);

            let ret = b.ins().iconst(types::I32, RET_CONTINUE as i64);
            b.ins().return_(&[ret]);
            b.finalize();
        }

        let id = self
            .module
            .declare_anonymous_function(&self.ctx.func.signature)
            .ok()?;
        self.module.define_function(id, &mut self.ctx).ok()?;
        self.module.clear_context(&mut self.ctx);
        self.module.finalize_definitions().ok()?;
        let code = self.module.get_finalized_function(id);
        // Safety: the module owns this code for its lifetime; the signature
        // matches CompiledFn (SystemV/extern "C" on the host).
        Some(unsafe { std::mem::transmute::<*const u8, CompiledFn>(code) })
    }
}

#[inline]
fn off(o: usize) -> i32 {
    o as i32
}

#[inline]
fn gpr_off(r: usize) -> i32 {
    (offset_of!(PpcContext, gpr) + r * 8) as i32
}

#[inline]
fn load_gpr(b: &mut FunctionBuilder, ctxp: Value, r: usize) -> Value {
    b.ins().load(types::I64, MemFlags::trusted(), ctxp, gpr_off(r))
}

#[inline]
fn store_gpr(b: &mut FunctionBuilder, ctxp: Value, r: usize, v: Value) {
    b.ins().store(MemFlags::trusted(), v, ctxp, gpr_off(r));
}

#[inline]
fn bump_u64(b: &mut FunctionBuilder, ctxp: Value, offset: usize, by: i64) {
    let cur = b.ins().load(types::I64, MemFlags::trusted(), ctxp, off(offset));
    let next = b.ins().iadd_imm(cur, by);
    b.ins().store(MemFlags::trusted(), next, ctxp, off(offset));
}

/// Emit IR for one covered instruction. Must mirror [`covered`] and the exact
/// semantics of the interpreter's `execute()` (validated by the diff harness).
fn emit_op(b: &mut FunctionBuilder, ctxp: Value, instr: &DecodedInstr) {
    match instr.opcode {
        // rd = (ra==0 ? 0 : gpr[ra]) + EXTS(simm)      [64-bit]
        PpcOpcode::addi => {
            let ra = instr.ra();
            let rav = if ra == 0 {
                b.ins().iconst(types::I64, 0)
            } else {
                load_gpr(b, ctxp, ra)
            };
            let res = b.ins().iadd_imm(rav, instr.simm16() as i64);
            store_gpr(b, ctxp, instr.rd(), res);
        }
        // rd = (ra==0 ? 0 : gpr[ra]) + (EXTS(simm) << 16)   [64-bit]
        PpcOpcode::addis => {
            let ra = instr.ra();
            let rav = if ra == 0 {
                b.ins().iconst(types::I64, 0)
            } else {
                load_gpr(b, ctxp, ra)
            };
            let res = b.ins().iadd_imm(rav, (instr.simm16() as i64) << 16);
            store_gpr(b, ctxp, instr.rd(), res);
        }
        // covered() gates this — unreachable for uncovered opcodes.
        _ => unreachable!("emit_op called on uncovered opcode {:?}", instr.opcode),
    }
}

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

    /// Mock memory that returns the same instruction word at every address.
    struct WordMem(u32);
    impl MemoryAccess for WordMem {
        fn read_u8(&self, _a: u32) -> u8 {
            0
        }
        fn read_u16(&self, _a: u32) -> u16 {
            0
        }
        fn read_u32(&self, _a: u32) -> u32 {
            self.0
        }
        fn read_u64(&self, _a: u32) -> u64 {
            ((self.0 as u64) << 32) | self.0 as u64
        }
        fn write_u8(&self, _a: u32, _v: u8) {}
        fn write_u16(&self, _a: u32, _v: u16) {}
        fn write_u32(&self, _a: u32, _v: u32) {}
        fn write_u64(&self, _a: u32, _v: u64) {}
        fn translate(&self, _a: u32) -> Option<*const u8> {
            None
        }
        fn translate_mut(&self, _a: u32) -> Option<*mut u8> {
            None
        }
    }

    // addi rD, rA, SIMM  =  (14<<26)|(rD<<21)|(rA<<16)|(SIMM & 0xFFFF)
    fn addi(rd: u32, ra: u32, simm: u16) -> u32 {
        (14 << 26) | (rd << 21) | (ra << 16) | simm as u32
    }

    #[test]
    fn jit_matches_interpreter_addi_block() {
        // A block of `addi r3, r3, 1` — 32 instrs (hits MAX_BLOCK_INSTRS, no
        // terminator), fully covered.
        let mem = WordMem(addi(3, 3, 1));
        let mut bc = BlockCache::new();
        let block = bc.lookup_or_build(0x1000, &mem).clone();
        assert!(block_covered(&block), "addi block should be covered");

        // Interpreter reference.
        let mut ref_ctx = PpcContext::new();
        ref_ctx.pc = 0x1000;
        let _ = crate::interpreter::step_block(&mut ref_ctx, &mem, &block);

        // JIT candidate.
        let mut jit = Jit::new().expect("jit init");
        let f = jit.compile(&block).expect("compile");
        let mut jit_ctx = PpcContext::new();
        jit_ctx.pc = 0x1000;
        let env = MemEnv { mem: &mem };
        let r = f(&mut jit_ctx as *mut _, &env as *const _);

        assert_eq!(r, RET_CONTINUE);
        assert_eq!(jit_ctx.gpr[3], ref_ctx.gpr[3], "r3 mismatch");
        assert_eq!(jit_ctx.pc, ref_ctx.pc, "pc mismatch");
        assert_eq!(jit_ctx.cycle_count, ref_ctx.cycle_count, "cycle mismatch");
        assert_eq!(jit_ctx.timebase, ref_ctx.timebase, "timebase mismatch");
    }
}