//! 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::condcodes::IntCC;
use cranelift_codegen::ir::{types, AbiParam, FuncRef, InstBuilder, MemFlags, Signature, Value};
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_jit::{JITBuilder, JITModule};
use cranelift_module::{default_libcall_names, FuncId, Linkage, 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;

// ---- memory trampolines ---------------------------------------------------
//
// Compiled loads/stores can't call the `&dyn MemoryAccess` trait method through
// its vtable from Cranelift IR, so they call these `extern "C"` shims, which do
// the dynamic dispatch in Rust. Routing through `MemoryAccess` keeps MMIO
// dispatch, mem-watch, `page_version`, and `mmio_access_count` bit-identical to
// the interpreter. Loads return the value zero-extended to u32/u64; stores take
// the value in the low bits of a u32 (or a full u64). Stores also receive the
// `PpcContext` pointer so they can replicate the interpreter's reservation
// invalidation (`stwcx` peers must observe an ordinary store to a reserved line).

extern "C" fn xj_read8(env: *const MemEnv, addr: u32) -> u32 {
    unsafe { (*env).mem.read_u8(addr) as u32 }
}
extern "C" fn xj_read16(env: *const MemEnv, addr: u32) -> u32 {
    unsafe { (*env).mem.read_u16(addr) as u32 }
}
extern "C" fn xj_read32(env: *const MemEnv, addr: u32) -> u32 {
    unsafe { (*env).mem.read_u32(addr) }
}
extern "C" fn xj_read64(env: *const MemEnv, addr: u32) -> u64 {
    unsafe { (*env).mem.read_u64(addr) }
}

/// Mirror the interpreter's pre-store reservation invalidation exactly.
#[inline]
fn store_reservation_kick(ctxp: *const PpcContext, addr: u32) {
    let ctx = unsafe { &*ctxp };
    if let Some(t) = ctx.reservation_table.as_ref().filter(|t| t.is_enabled()) {
        if t.has_active_reservers() {
            t.invalidate_for_write(addr);
        }
    }
}

extern "C" fn xj_write8(ctxp: *const PpcContext, env: *const MemEnv, addr: u32, val: u32) {
    store_reservation_kick(ctxp, addr);
    unsafe { (*env).mem.write_u8(addr, val as u8) }
}
extern "C" fn xj_write16(ctxp: *const PpcContext, env: *const MemEnv, addr: u32, val: u32) {
    store_reservation_kick(ctxp, addr);
    unsafe { (*env).mem.write_u16(addr, val as u16) }
}
extern "C" fn xj_write32(ctxp: *const PpcContext, env: *const MemEnv, addr: u32, val: u32) {
    store_reservation_kick(ctxp, addr);
    unsafe { (*env).mem.write_u32(addr, val) }
}
extern "C" fn xj_write64(ctxp: *const PpcContext, env: *const MemEnv, addr: u32, val: u64) {
    store_reservation_kick(ctxp, addr);
    unsafe { (*env).mem.write_u64(addr, val) }
}

/// `FuncRef`s for the eight trampolines, declared into the function being built.
/// Copied into each block's IR so load/store arms can emit calls.
#[derive(Clone, Copy)]
struct Trampolines {
    read8: FuncRef,
    read16: FuncRef,
    read32: FuncRef,
    read64: FuncRef,
    write8: FuncRef,
    write16: FuncRef,
    write32: FuncRef,
    write64: FuncRef,
}

/// Per-compile emit context: the two entry-block pointer params plus the
/// trampoline `FuncRef`s. `ctxp` is `*mut PpcContext`, `memenv` is `*const
/// MemEnv`, both as IR `Value`s.
struct EmitCtx {
    ctxp: Value,
    memenv: Value,
    tr: Trampolines,
}

// ---- compiled-block cache -------------------------------------------------

/// Direct-mapped compiled-block cache — same slot geometry and `(start_pc,
/// page_version)` keying as [`crate::block_cache::BlockCache`], so a
/// self-modifying / DMA'd code page invalidates compiled code the same way it
/// invalidates the decoded block (both re-key on the bumped `page_version`).
const JIT_CACHE_SIZE: usize = 1 << 16;
const JIT_CACHE_MASK: u32 = (JIT_CACHE_SIZE - 1) as u32;

/// One compiled-cache slot. `func == None` records "this block is **uncovered**
/// (or failed to compile) — don't retry", so an uncovered block is compile-
/// attempted at most once per `(pc, page_version)` instead of every visit.
struct CompiledSlot {
    start_pc: u32,
    page_version: u64,
    func: Option<CompiledFn>,
}

/// A [`Jit`] plus a direct-mapped cache of its compiled blocks. One instance
/// lives per worker slot (like `BlockCache`); the owned [`Jit`] module keeps
/// every [`CompiledFn`] it hands out valid for the cache's lifetime.
pub struct JitCache {
    /// `None` only if Cranelift failed to initialise on this host — then every
    /// lookup misses and the caller interprets (no JIT available).
    jit: Option<Jit>,
    slots: Box<[Option<CompiledSlot>]>,
}

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

impl JitCache {
    pub fn new() -> Self {
        let mut slots: Vec<Option<CompiledSlot>> = Vec::with_capacity(JIT_CACHE_SIZE);
        slots.resize_with(JIT_CACHE_SIZE, || None);
        Self { jit: Jit::new().ok(), slots: slots.into_boxed_slice() }
    }

    /// Return the compiled entry point for `block`, compiling it on the first
    /// visit. `None` means the block is uncovered (or compilation failed) and
    /// the caller must interpret it. The `None` verdict is cached per
    /// `(start_pc, page_version)` so uncovered blocks aren't re-attempted.
    pub fn get_or_compile(&mut self, block: &DecodedBlock) -> Option<CompiledFn> {
        let idx = ((block.start_pc >> 2) & JIT_CACHE_MASK) as usize;
        if let Some(slot) = &self.slots[idx] {
            if slot.start_pc == block.start_pc && slot.page_version == block.page_version {
                return slot.func;
            }
        }
        let func = self.jit.as_mut().and_then(|j| j.compile(block));
        self.slots[idx] = Some(CompiledSlot {
            start_pc: block.start_pc,
            page_version: block.page_version,
            func,
        });
        func
    }
}

/// 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.
///
/// Covered so far: `addi`/`addis` (integer immediate add) and the three direct
/// branch terminators `bx`/`bcx`/`bclrx`. `bcctrx` is deliberately excluded —
/// it targets an indirect address and runs the `dispatch_rec` diagnostic hook,
/// which the native path would silently skip. Every branch computes its target
/// from immediates (or the live `lr` for `bclrx`) exactly as the interpreter.
pub fn covered(instr: &DecodedInstr) -> bool {
    use PpcOpcode::*;
    match instr.opcode {
        addi | addis => true,
        // Reg-reg logical + immediate logical + rotate + add/sub + compare.
        // `orx` 0x7FFFFB78 is the db16cyc spin hint (returns Yield) — leave to
        // the interpreter. `addx`/`subfx` with OE set update XER OV/SO via the
        // overflow path — not lowered yet, so only cover OE=0.
        orx => instr.raw != 0x7FFF_FB78,
        andx | xorx | norx | nandx | andcx | orcx => true,
        ori | oris | xori | xoris | andix | andisx => true,
        addx | subfx => !instr.oe(),
        rlwinmx => true,
        cmp | cmpl | cmpi | cmpli => true,
        // Integer loads/stores via memory trampolines (zero-extended loads and
        // the non-update forms only). Update (`u`) forms and algebraic
        // (sign-extending) loads are not lowered yet.
        lbz | lbzx | lhz | lhzx | lwz | lwzx | ld | ldx => true,
        stb | stbx | sth | sthx | stw | stwx | std | stdx => true,
        // Branch terminators.
        bx | bcx | bclrx => true,
        _ => false,
    }
}

/// Does this opcode write `pc` itself (a branch), so the block epilogue must
/// **not** overwrite `pc` with `end_pc`? Only ever true for the terminator.
fn writes_pc(instr: &DecodedInstr) -> bool {
    matches!(instr.opcode, PpcOpcode::bx | PpcOpcode::bcx | PpcOpcode::bclrx)
}

/// 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,
    /// Imported trampoline function ids (declared once; re-referenced into each
    /// compiled function via `declare_func_in_func`).
    tramp_ids: TrampIds,
}

/// `FuncId`s of the imported memory trampolines.
struct TrampIds {
    read8: FuncId,
    read16: FuncId,
    read32: FuncId,
    read64: FuncId,
    write8: FuncId,
    write16: FuncId,
    write32: FuncId,
    write64: FuncId,
}

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 mut builder = JITBuilder::new(default_libcall_names()).map_err(|e| e.to_string())?;
        // Register the trampoline addresses so `Linkage::Import` symbols resolve.
        builder.symbol("xj_read8", xj_read8 as *const u8);
        builder.symbol("xj_read16", xj_read16 as *const u8);
        builder.symbol("xj_read32", xj_read32 as *const u8);
        builder.symbol("xj_read64", xj_read64 as *const u8);
        builder.symbol("xj_write8", xj_write8 as *const u8);
        builder.symbol("xj_write16", xj_write16 as *const u8);
        builder.symbol("xj_write32", xj_write32 as *const u8);
        builder.symbol("xj_write64", xj_write64 as *const u8);
        let mut module = JITModule::new(builder);

        let ptr = module.target_config().pointer_type();
        let cc = module.target_config().default_call_conv;
        // read(env: ptr, addr: i32) -> {i32|i64}
        let load_sig = |ret| {
            let mut s = Signature::new(cc);
            s.params.push(AbiParam::new(ptr));
            s.params.push(AbiParam::new(types::I32));
            s.returns.push(AbiParam::new(ret));
            s
        };
        // write(ctx: ptr, env: ptr, addr: i32, val: {i32|i64})
        let store_sig = |valty| {
            let mut s = Signature::new(cc);
            s.params.push(AbiParam::new(ptr));
            s.params.push(AbiParam::new(ptr));
            s.params.push(AbiParam::new(types::I32));
            s.params.push(AbiParam::new(valty));
            s
        };
        let mut declare = |name: &str, sig: &Signature| -> Result<FuncId, String> {
            module.declare_function(name, Linkage::Import, sig).map_err(|e| e.to_string())
        };
        let tramp_ids = TrampIds {
            read8: declare("xj_read8", &load_sig(types::I32))?,
            read16: declare("xj_read16", &load_sig(types::I32))?,
            read32: declare("xj_read32", &load_sig(types::I32))?,
            read64: declare("xj_read64", &load_sig(types::I64))?,
            write8: declare("xj_write8", &store_sig(types::I32))?,
            write16: declare("xj_write16", &store_sig(types::I32))?,
            write32: declare("xj_write32", &store_sig(types::I32))?,
            write64: declare("xj_write64", &store_sig(types::I64))?,
        };

        let ctx = module.make_context();
        Ok(Self { module, ctx, fbctx: FunctionBuilderContext::new(), tramp_ids })
    }

    /// 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));

        // Import the trampolines into this function (before the builder borrows
        // `ctx.func`). `declare_func_in_func` borrows the module immutably and
        // `ctx.func` mutably — disjoint fields, so no borrow conflict.
        let tr = Trampolines {
            read8: self.module.declare_func_in_func(self.tramp_ids.read8, &mut self.ctx.func),
            read16: self.module.declare_func_in_func(self.tramp_ids.read16, &mut self.ctx.func),
            read32: self.module.declare_func_in_func(self.tramp_ids.read32, &mut self.ctx.func),
            read64: self.module.declare_func_in_func(self.tramp_ids.read64, &mut self.ctx.func),
            write8: self.module.declare_func_in_func(self.tramp_ids.write8, &mut self.ctx.func),
            write16: self.module.declare_func_in_func(self.tramp_ids.write16, &mut self.ctx.func),
            write32: self.module.declare_func_in_func(self.tramp_ids.write32, &mut self.ctx.func),
            write64: self.module.declare_func_in_func(self.tramp_ids.write64, &mut self.ctx.func),
        };

        {
            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 ec = EmitCtx {
                ctxp: b.block_params(entry)[0],
                memenv: b.block_params(entry)[1],
                tr,
            };
            let ctxp = ec.ctxp;

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

            // A branch terminator writes `pc` itself (to its computed target);
            // otherwise the block is straight-line and `pc` advances to
            // `end_pc`. `writes_pc` is only ever true for the last instruction
            // (branches terminate the block), so testing it is sufficient.
            if !block.instrs.last().is_some_and(writes_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));
}

#[inline]
fn store_pc(b: &mut FunctionBuilder, ctxp: Value, val: Value) {
    b.ins().store(MemFlags::trusted(), val, ctxp, off(offset_of!(PpcContext, pc)));
}

#[inline]
fn store_lr(b: &mut FunctionBuilder, ctxp: Value, val_u64: u64) {
    let v = b.ins().iconst(types::I64, val_u64 as i64);
    b.ins().store(MemFlags::trusted(), v, ctxp, off(offset_of!(PpcContext, lr)));
}

/// Low 32 bits of GPR `r` as an `I32`.
#[inline]
fn gpr32(b: &mut FunctionBuilder, ctxp: Value, r: usize) -> Value {
    let v = load_gpr(b, ctxp, r);
    b.ins().ireduce(types::I32, v)
}

/// Store an `I32` result into GPR `ra`, zero-extended to 64 bits — the width
/// contract of the `*cx`/`nor`/`nand`/`rlwinm` ops (which zero the upper half).
#[inline]
fn store_gpr32z(b: &mut FunctionBuilder, ctxp: Value, ra: usize, v32: Value) {
    let z = b.ins().uextend(types::I64, v32);
    store_gpr(b, ctxp, ra, z);
}

/// PPC `rlwnm`-family mask. Mirrors the interpreter's `rlw_mask`; `mb`/`me` are
/// instruction immediates so the whole mask is a compile-time constant.
fn rlw_mask(mb: u32, me: u32) -> u32 {
    if mb <= me {
        (u32::MAX >> mb) & (u32::MAX << (31 - me))
    } else {
        (u32::MAX >> mb) | (u32::MAX << (31 - me))
    }
}

/// Store CR field `field` from three precomputed `I8` (0/1) predicates plus the
/// `so` bit copied from `xer_so`. `CrField` is `{lt@0, gt@1, eq@2, so@3}`, one
/// byte each, at `cr + field*4`.
fn emit_store_cr(
    b: &mut FunctionBuilder,
    ctxp: Value,
    field: usize,
    lt: Value,
    gt: Value,
    eq: Value,
) {
    let base = (offset_of!(PpcContext, cr) + field * 4) as i32;
    b.ins().store(MemFlags::trusted(), lt, ctxp, base);
    b.ins().store(MemFlags::trusted(), gt, ctxp, base + 1);
    b.ins().store(MemFlags::trusted(), eq, ctxp, base + 2);
    let so_raw = b.ins().load(
        types::I8,
        MemFlags::trusted(),
        ctxp,
        off(offset_of!(PpcContext, xer_so)),
    );
    let so = b.ins().icmp_imm(IntCC::NotEqual, so_raw, 0);
    b.ins().store(MemFlags::trusted(), so, ctxp, base + 3);
}

/// Emit `update_cr_signed(0, val as i32 as i64)` — the Rc-form CR0 update: a
/// signed comparison of the 32-bit result against zero.
fn emit_cr0_signed32(b: &mut FunctionBuilder, ctxp: Value, val32: Value) {
    let lt = b.ins().icmp_imm(IntCC::SignedLessThan, val32, 0);
    let gt = b.ins().icmp_imm(IntCC::SignedGreaterThan, val32, 0);
    let eq = b.ins().icmp_imm(IntCC::Equal, val32, 0);
    emit_store_cr(b, ctxp, 0, lt, gt, eq);
}

/// Load the single CR bit `bi` (0-31) as an `I8` boolean (0/1). CR fields are
/// `#[repr] struct CrField { lt, gt, eq, so }` — one byte each — so bit
/// `bi = field*4 + sub` is the byte at `cr + field*4 + sub`.
#[inline]
fn emit_cr_bit(b: &mut FunctionBuilder, ctxp: Value, bi: u32) -> Value {
    let field = (bi / 4) as usize;
    let sub = (bi % 4) as usize;
    let offset = (offset_of!(PpcContext, cr) + field * 4 + sub) as i32;
    b.ins().load(types::I8, MemFlags::trusted(), ctxp, offset)
}

/// Emit the shared `bcx`/`bclrx` "branch taken?" predicate, mirroring the
/// interpreter exactly: optionally decrement CTR, then `ctr_ok && cond_ok`.
/// Returns an `I8` (0/1). BO/BI are immediates, so the constant sub-cases
/// (`branch always`, `no CTR`, `ignore CR`) fold away at compile time.
fn emit_branch_taken(b: &mut FunctionBuilder, ctxp: Value, bo: u32, bi: u32) -> Value {
    let ctr_off = off(offset_of!(PpcContext, ctr));

    // ctr_ok. BO2 (bo & 0b00100): 1 = don't test CTR. Else decrement CTR and
    // test (ctr!=0) possibly inverted by BO3 (bo & 0b00010).
    let ctr_ok = if bo & 0b00100 != 0 {
        b.ins().iconst(types::I8, 1)
    } else {
        let cur = b.ins().load(types::I64, MemFlags::trusted(), ctxp, ctr_off);
        let dec = b.ins().iadd_imm(cur, -1);
        b.ins().store(MemFlags::trusted(), dec, ctxp, ctr_off);
        let lo = b.ins().ireduce(types::I32, dec);
        // interpreter: (ctr as u32 != 0) ^ (bo&2 != 0)
        let cc = if bo & 0b00010 != 0 { IntCC::Equal } else { IntCC::NotEqual };
        b.ins().icmp_imm(cc, lo, 0)
    };

    // cond_ok. BO0 (bo & 0b10000): 1 = don't test CR. Else CR bit BI must equal
    // BO1 (bo & 0b01000).
    let cond_ok = if bo & 0b10000 != 0 {
        b.ins().iconst(types::I8, 1)
    } else {
        let crbit = emit_cr_bit(b, ctxp, bi);
        if bo & 0b01000 != 0 {
            crbit
        } else {
            b.ins().icmp_imm(IntCC::Equal, crbit, 0)
        }
    };

    b.ins().band(ctr_ok, cond_ok)
}

/// Effective address for a D-form load/store: `(ra==0 ? 0 : gpr[ra]) + disp`,
/// truncated to 32 bits. Returns an `I32`.
fn ea_d(b: &mut FunctionBuilder, ec: &EmitCtx, ra: usize, disp: i32) -> Value {
    let base = if ra == 0 {
        b.ins().iconst(types::I64, 0)
    } else {
        load_gpr(b, ec.ctxp, ra)
    };
    let ea64 = b.ins().iadd_imm(base, disp as i64);
    b.ins().ireduce(types::I32, ea64)
}

/// Effective address for an X-form (indexed) load/store: `(ra==0 ? 0 : gpr[ra])
/// + gpr[rb]`, truncated to 32 bits. Returns an `I32`.
fn ea_x(b: &mut FunctionBuilder, ec: &EmitCtx, ra: usize, rb: usize) -> Value {
    let base = if ra == 0 {
        b.ins().iconst(types::I64, 0)
    } else {
        load_gpr(b, ec.ctxp, ra)
    };
    let rbv = load_gpr(b, ec.ctxp, rb);
    let ea64 = b.ins().iadd(base, rbv);
    b.ins().ireduce(types::I32, ea64)
}

/// Emit a load trampoline call `fref(memenv, ea)` and return its result value.
fn call_load(b: &mut FunctionBuilder, ec: &EmitCtx, fref: FuncRef, ea: Value) -> Value {
    let call = b.ins().call(fref, &[ec.memenv, ea]);
    b.inst_results(call)[0]
}

/// Emit a store trampoline call `fref(ctxp, memenv, ea, val)`.
fn call_store(b: &mut FunctionBuilder, ec: &EmitCtx, fref: FuncRef, ea: Value, val: Value) {
    b.ins().call(fref, &[ec.ctxp, ec.memenv, ea, val]);
}

/// 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, ec: &EmitCtx, instr: &DecodedInstr) {
    let ctxp = ec.ctxp;
    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);
        }

        // ---- reg-reg add/sub (OE=0). Full 64-bit result; CR0 on low 32. ----
        PpcOpcode::addx => {
            let ra = load_gpr(b, ctxp, instr.ra());
            let rb = load_gpr(b, ctxp, instr.rb());
            let res = b.ins().iadd(ra, rb);
            store_gpr(b, ctxp, instr.rd(), res);
            if instr.rc_bit() {
                let lo = b.ins().ireduce(types::I32, res);
                emit_cr0_signed32(b, ctxp, lo);
            }
        }
        PpcOpcode::subfx => {
            // rd = rb - ra
            let ra = load_gpr(b, ctxp, instr.ra());
            let rb = load_gpr(b, ctxp, instr.rb());
            let res = b.ins().isub(rb, ra);
            store_gpr(b, ctxp, instr.rd(), res);
            if instr.rc_bit() {
                let lo = b.ins().ireduce(types::I32, res);
                emit_cr0_signed32(b, ctxp, lo);
            }
        }

        // ---- reg-reg logical preserving all 64 bits; CR0 on low 32 if Rc. ----
        PpcOpcode::orx => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let rb = load_gpr(b, ctxp, instr.rb());
            let res = b.ins().bor(rs, rb);
            store_gpr(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                let lo = b.ins().ireduce(types::I32, res);
                emit_cr0_signed32(b, ctxp, lo);
            }
        }
        PpcOpcode::andx => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let rb = load_gpr(b, ctxp, instr.rb());
            let res = b.ins().band(rs, rb);
            store_gpr(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                let lo = b.ins().ireduce(types::I32, res);
                emit_cr0_signed32(b, ctxp, lo);
            }
        }
        PpcOpcode::xorx => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let rb = load_gpr(b, ctxp, instr.rb());
            let res = b.ins().bxor(rs, rb);
            store_gpr(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                let lo = b.ins().ireduce(types::I32, res);
                emit_cr0_signed32(b, ctxp, lo);
            }
        }

        // ---- reg-reg logical operating in u32 (zeroes upper 32); CR0 low 32. ----
        PpcOpcode::norx => {
            let rs = gpr32(b, ctxp, instr.rs());
            let rb = gpr32(b, ctxp, instr.rb());
            let t = b.ins().bor(rs, rb);
            let res = b.ins().bnot(t);
            store_gpr32z(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                emit_cr0_signed32(b, ctxp, res);
            }
        }
        PpcOpcode::nandx => {
            let rs = gpr32(b, ctxp, instr.rs());
            let rb = gpr32(b, ctxp, instr.rb());
            let t = b.ins().band(rs, rb);
            let res = b.ins().bnot(t);
            store_gpr32z(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                emit_cr0_signed32(b, ctxp, res);
            }
        }
        PpcOpcode::andcx => {
            let rs = gpr32(b, ctxp, instr.rs());
            let rb = gpr32(b, ctxp, instr.rb());
            let nrb = b.ins().bnot(rb);
            let res = b.ins().band(rs, nrb);
            store_gpr32z(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                emit_cr0_signed32(b, ctxp, res);
            }
        }
        PpcOpcode::orcx => {
            let rs = gpr32(b, ctxp, instr.rs());
            let rb = gpr32(b, ctxp, instr.rb());
            let nrb = b.ins().bnot(rb);
            let res = b.ins().bor(rs, nrb);
            store_gpr32z(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                emit_cr0_signed32(b, ctxp, res);
            }
        }

        // ---- immediate logical. 64-bit result. ori/oris/xori/xoris: no CR. ----
        PpcOpcode::ori => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let res = b.ins().bor_imm(rs, instr.uimm16() as i64);
            store_gpr(b, ctxp, instr.ra(), res);
        }
        PpcOpcode::oris => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let res = b.ins().bor_imm(rs, (instr.uimm16() as i64) << 16);
            store_gpr(b, ctxp, instr.ra(), res);
        }
        PpcOpcode::xori => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let res = b.ins().bxor_imm(rs, instr.uimm16() as i64);
            store_gpr(b, ctxp, instr.ra(), res);
        }
        PpcOpcode::xoris => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let res = b.ins().bxor_imm(rs, (instr.uimm16() as i64) << 16);
            store_gpr(b, ctxp, instr.ra(), res);
        }
        // andi./andis. always update CR0 (the `.` forms).
        PpcOpcode::andix => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let res = b.ins().band_imm(rs, instr.uimm16() as i64);
            store_gpr(b, ctxp, instr.ra(), res);
            let lo = b.ins().ireduce(types::I32, res);
            emit_cr0_signed32(b, ctxp, lo);
        }
        PpcOpcode::andisx => {
            let rs = load_gpr(b, ctxp, instr.rs());
            let res = b.ins().band_imm(rs, (instr.uimm16() as i64) << 16);
            store_gpr(b, ctxp, instr.ra(), res);
            let lo = b.ins().ireduce(types::I32, res);
            emit_cr0_signed32(b, ctxp, lo);
        }

        // ---- rotate-left word immediate then AND mask; zeroes upper 32. ----
        PpcOpcode::rlwinmx => {
            let rs = gpr32(b, ctxp, instr.rs());
            let shv = b.ins().iconst(types::I32, instr.sh() as i64);
            let rot = b.ins().rotl(rs, shv);
            let mask = rlw_mask(instr.mb(), instr.me());
            let res = b.ins().band_imm(rot, mask as i64);
            store_gpr32z(b, ctxp, instr.ra(), res);
            if instr.rc_bit() {
                emit_cr0_signed32(b, ctxp, res);
            }
        }

        // ---- compares. Write CR field crfd() unconditionally. ----
        PpcOpcode::cmp => {
            let (lt, gt, eq) = if instr.l() {
                let ra = load_gpr(b, ctxp, instr.ra());
                let rb = load_gpr(b, ctxp, instr.rb());
                (
                    b.ins().icmp(IntCC::SignedLessThan, ra, rb),
                    b.ins().icmp(IntCC::SignedGreaterThan, ra, rb),
                    b.ins().icmp(IntCC::Equal, ra, rb),
                )
            } else {
                let ra = gpr32(b, ctxp, instr.ra());
                let rb = gpr32(b, ctxp, instr.rb());
                (
                    b.ins().icmp(IntCC::SignedLessThan, ra, rb),
                    b.ins().icmp(IntCC::SignedGreaterThan, ra, rb),
                    b.ins().icmp(IntCC::Equal, ra, rb),
                )
            };
            emit_store_cr(b, ctxp, instr.crfd(), lt, gt, eq);
        }
        PpcOpcode::cmpl => {
            let (lt, gt, eq) = if instr.l() {
                let ra = load_gpr(b, ctxp, instr.ra());
                let rb = load_gpr(b, ctxp, instr.rb());
                (
                    b.ins().icmp(IntCC::UnsignedLessThan, ra, rb),
                    b.ins().icmp(IntCC::UnsignedGreaterThan, ra, rb),
                    b.ins().icmp(IntCC::Equal, ra, rb),
                )
            } else {
                let ra = gpr32(b, ctxp, instr.ra());
                let rb = gpr32(b, ctxp, instr.rb());
                (
                    b.ins().icmp(IntCC::UnsignedLessThan, ra, rb),
                    b.ins().icmp(IntCC::UnsignedGreaterThan, ra, rb),
                    b.ins().icmp(IntCC::Equal, ra, rb),
                )
            };
            emit_store_cr(b, ctxp, instr.crfd(), lt, gt, eq);
        }
        PpcOpcode::cmpi => {
            let imm = instr.simm16() as i64; // sign-extended
            let (lt, gt, eq) = if instr.l() {
                let ra = load_gpr(b, ctxp, instr.ra());
                (
                    b.ins().icmp_imm(IntCC::SignedLessThan, ra, imm),
                    b.ins().icmp_imm(IntCC::SignedGreaterThan, ra, imm),
                    b.ins().icmp_imm(IntCC::Equal, ra, imm),
                )
            } else {
                let ra = gpr32(b, ctxp, instr.ra());
                // 32-bit signed compare against sign-extended SIMM.
                let imm32 = instr.simm16() as i32 as i64;
                (
                    b.ins().icmp_imm(IntCC::SignedLessThan, ra, imm32),
                    b.ins().icmp_imm(IntCC::SignedGreaterThan, ra, imm32),
                    b.ins().icmp_imm(IntCC::Equal, ra, imm32),
                )
            };
            emit_store_cr(b, ctxp, instr.crfd(), lt, gt, eq);
        }
        PpcOpcode::cmpli => {
            let imm = instr.uimm16() as i64; // zero-extended
            let (lt, gt, eq) = if instr.l() {
                let ra = load_gpr(b, ctxp, instr.ra());
                (
                    b.ins().icmp_imm(IntCC::UnsignedLessThan, ra, imm),
                    b.ins().icmp_imm(IntCC::UnsignedGreaterThan, ra, imm),
                    b.ins().icmp_imm(IntCC::Equal, ra, imm),
                )
            } else {
                let ra = gpr32(b, ctxp, instr.ra());
                (
                    b.ins().icmp_imm(IntCC::UnsignedLessThan, ra, imm),
                    b.ins().icmp_imm(IntCC::UnsignedGreaterThan, ra, imm),
                    b.ins().icmp_imm(IntCC::Equal, ra, imm),
                )
            };
            emit_store_cr(b, ctxp, instr.crfd(), lt, gt, eq);
        }

        // ---- integer loads (zero-extended into the 64-bit GPR). ----
        PpcOpcode::lbz => {
            let ea = ea_d(b, ec, instr.ra(), instr.d());
            let v = call_load(b, ec, ec.tr.read8, ea);
            let z = b.ins().uextend(types::I64, v);
            store_gpr(b, ctxp, instr.rd(), z);
        }
        PpcOpcode::lbzx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = call_load(b, ec, ec.tr.read8, ea);
            let z = b.ins().uextend(types::I64, v);
            store_gpr(b, ctxp, instr.rd(), z);
        }
        PpcOpcode::lhz => {
            let ea = ea_d(b, ec, instr.ra(), instr.d());
            let v = call_load(b, ec, ec.tr.read16, ea);
            let z = b.ins().uextend(types::I64, v);
            store_gpr(b, ctxp, instr.rd(), z);
        }
        PpcOpcode::lhzx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = call_load(b, ec, ec.tr.read16, ea);
            let z = b.ins().uextend(types::I64, v);
            store_gpr(b, ctxp, instr.rd(), z);
        }
        PpcOpcode::lwz => {
            let ea = ea_d(b, ec, instr.ra(), instr.d());
            let v = call_load(b, ec, ec.tr.read32, ea);
            let z = b.ins().uextend(types::I64, v);
            store_gpr(b, ctxp, instr.rd(), z);
        }
        PpcOpcode::lwzx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = call_load(b, ec, ec.tr.read32, ea);
            let z = b.ins().uextend(types::I64, v);
            store_gpr(b, ctxp, instr.rd(), z);
        }
        PpcOpcode::ld => {
            let ea = ea_d(b, ec, instr.ra(), instr.ds());
            let v = call_load(b, ec, ec.tr.read64, ea);
            store_gpr(b, ctxp, instr.rd(), v);
        }
        PpcOpcode::ldx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = call_load(b, ec, ec.tr.read64, ea);
            store_gpr(b, ctxp, instr.rd(), v);
        }

        // ---- integer stores. Value is the low bits of gpr[rs]. ----
        PpcOpcode::stb => {
            let ea = ea_d(b, ec, instr.ra(), instr.d());
            let v = gpr32(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write8, ea, v);
        }
        PpcOpcode::stbx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = gpr32(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write8, ea, v);
        }
        PpcOpcode::sth => {
            let ea = ea_d(b, ec, instr.ra(), instr.d());
            let v = gpr32(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write16, ea, v);
        }
        PpcOpcode::sthx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = gpr32(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write16, ea, v);
        }
        PpcOpcode::stw => {
            let ea = ea_d(b, ec, instr.ra(), instr.d());
            let v = gpr32(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write32, ea, v);
        }
        PpcOpcode::stwx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = gpr32(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write32, ea, v);
        }
        PpcOpcode::std => {
            let ea = ea_d(b, ec, instr.ra(), instr.ds());
            let v = load_gpr(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write64, ea, v);
        }
        PpcOpcode::stdx => {
            let ea = ea_x(b, ec, instr.ra(), instr.rb());
            let v = load_gpr(b, ctxp, instr.rs());
            call_store(b, ec, ec.tr.write64, ea, v);
        }

        // Unconditional branch. Target is an immediate; `pc` (= this instr's
        // address in the interpreter) is `instr.addr` at emit time.
        PpcOpcode::bx => {
            let target = if instr.aa() {
                instr.li() as u32
            } else {
                instr.addr.wrapping_add(instr.li() as u32)
            };
            if instr.lk() {
                store_lr(b, ctxp, instr.addr.wrapping_add(4) as u64);
            }
            let t = b.ins().iconst(types::I32, target as i64);
            store_pc(b, ctxp, t);
        }
        // Conditional branch (relative). CTR decrement / CR test via
        // `emit_branch_taken`; both targets are immediates.
        PpcOpcode::bcx => {
            let next = instr.addr.wrapping_add(4);
            let target = if instr.aa() {
                instr.bd() as u32
            } else {
                instr.addr.wrapping_add(instr.bd() as u32)
            };
            let taken = emit_branch_taken(b, ctxp, instr.bo(), instr.bi());
            let t = b.ins().iconst(types::I32, target as i64);
            let n = b.ins().iconst(types::I32, next as i64);
            let pc = b.ins().select(taken, t, n);
            store_pc(b, ctxp, pc);
            // interpreter sets LR = next unconditionally when LK, on both paths.
            if instr.lk() {
                store_lr(b, ctxp, next as u64);
            }
        }
        // Branch-conditional to LR. Target is the live `lr & !3`, read BEFORE a
        // LK link overwrites it.
        PpcOpcode::bclrx => {
            let next = instr.addr.wrapping_add(4);
            let lr_cur = b.ins().load(
                types::I64,
                MemFlags::trusted(),
                ctxp,
                off(offset_of!(PpcContext, lr)),
            );
            let lr32 = b.ins().ireduce(types::I32, lr_cur);
            let target = b.ins().band_imm(lr32, !3i64);
            let taken = emit_branch_taken(b, ctxp, instr.bo(), instr.bi());
            let n = b.ins().iconst(types::I32, next as i64);
            let pc = b.ins().select(taken, target, n);
            store_pc(b, ctxp, pc);
            if instr.lk() {
                store_lr(b, ctxp, next as u64);
            }
        }
        // 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");
    }
}