[iterate-4A] jit: cache MemEnv fast_mem + expose mmio_count (prep for block chaining)

Behavior-neutral prep on top of the inline-load parity milestone:

  * JitCache caches mem.fast_mem() (resolved once — the mapping is invariant for
    a run) so compiled blocks build their MemEnv without a virtual call per
    execution. MemEnv::from_fast(mem, Option<FastMem>) is the non-virtual
    constructor. Measured impact is within noise (96.5 -> 96.4 MIPS) but it
    removes redundant per-block work.
  * FastMem/MemEnv gain mmio_count: a pointer to GuestMemory's monotonic MMIO
    access counter. Unused for now; the upcoming superblock-chaining JIT will
    sample it across a block boundary to stop chaining on an MMIO touch
    (preserving the interpreter's fine-grained MMIO ordering).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xenia-cpu/src/jit.rs

@@ -33,7 +33,7 @@ use crate::block_cache::DecodedBlock;
 use crate::context::PpcContext;
 use crate::decoder::DecodedInstr;
 use crate::opcode::PpcOpcode;
-use xenia_memory::MemoryAccess;
+use xenia_memory::{FastMem, MemoryAccess};
 
 /// `StepResult` discriminants returned by compiled blocks. Kept in sync with
 /// [`crate::interpreter::StepResult`] (only `Continue` is produced by covered
@@ -56,6 +56,7 @@ pub struct MemEnv<'a> {
     pub page_table: *const u64,
     pub mmio_mask: u32,
     pub mmio_value: u32,
+    pub mmio_count: *const u64,
 }
 
 impl<'a> MemEnv<'a> {
@@ -63,13 +64,21 @@ impl<'a> MemEnv<'a> {
     /// exposes them ([`MemoryAccess::fast_mem`]). Memories that intercept
     /// accesses return `None`, leaving `membase` null (fast path disabled).
     pub fn new(mem: &'a dyn MemoryAccess) -> Self {
-        match mem.fast_mem() {
+        Self::from_fast(mem, mem.fast_mem())
+    }
+
+    /// Build the env from `mem` plus an already-resolved [`FastMem`] — lets the
+    /// caller cache `mem.fast_mem()` (a virtual call) instead of paying it per
+    /// block. `fm` must belong to `mem`.
+    pub fn from_fast(mem: &'a dyn MemoryAccess, fm: Option<FastMem>) -> Self {
+        match fm {
             Some(f) => Self {
                 mem,
                 membase: f.membase,
                 page_table: f.page_table,
                 mmio_mask: f.mmio_mask,
                 mmio_value: f.mmio_value,
+                mmio_count: f.mmio_count,
             },
             None => Self {
                 mem,
@@ -77,6 +86,7 @@ impl<'a> MemEnv<'a> {
                 page_table: std::ptr::null(),
                 mmio_mask: 0,
                 mmio_value: 1, // (addr & 0) == 1 is never true; membase gate wins anyway
+                mmio_count: std::ptr::null(),
             },
         }
     }
@@ -204,6 +214,11 @@ pub struct JitCache {
     /// lookup misses and the caller interprets (no JIT available).
     jit: Option<Jit>,
     slots: Box<[Option<CompiledSlot>]>,
+    /// Cached `mem.fast_mem()` — resolved once (the mapping is invariant for a
+    /// run) so compiled blocks don't pay a virtual call per execution to build
+    /// their `MemEnv`. `fast_mem_init` guards first use.
+    fast_mem: Option<FastMem>,
+    fast_mem_init: bool,
 }
 
 impl Default for JitCache {
@@ -216,7 +231,23 @@ 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() }
+        Self {
+            jit: Jit::new().ok(),
+            slots: slots.into_boxed_slice(),
+            fast_mem: None,
+            fast_mem_init: false,
+        }
+    }
+
+    /// Build a `MemEnv` for `mem`, resolving (and caching) its `fast_mem()` on
+    /// first use. `mem` is invariant across a run, so the cached value stays
+    /// correct (the overlay used under the diff harness always yields `None`).
+    pub fn mem_env<'a>(&mut self, mem: &'a dyn MemoryAccess) -> MemEnv<'a> {
+        if !self.fast_mem_init {
+            self.fast_mem = mem.fast_mem();
+            self.fast_mem_init = true;
+        }
+        MemEnv::from_fast(mem, self.fast_mem)
     }
 
     /// Return the compiled entry point for `block`, compiling it on the first

crates/xenia-cpu/src/recompiler.rs

@@ -150,7 +150,7 @@ pub fn run_block(
 ) -> StepResult {
     if let Some(f) = jit.get_or_compile(block) {
         JIT_COMPILED_RUN.fetch_add(1, Ordering::Relaxed);
-        let env = MemEnv::new(mem);
+        let env = jit.mem_env(mem);
         let raw = f(ctx as *mut PpcContext, &env as *const MemEnv);
         // Covered blocks are straight-line (branch/sc/trap/db16cyc are all
         // uncovered), so a compiled block always runs to completion and

crates/xenia-memory/src/access.rs

@@ -22,6 +22,11 @@ pub struct FastMem {
     /// condition for `addr` to be MMIO, so `!=` proves non-MMIO.
     pub mmio_mask: u32,
     pub mmio_value: u32,
+    /// Pointer to the monotonic MMIO-access counter (`AtomicU64` as `*const
+    /// u64`). A superblock JIT samples this across a block boundary to detect an
+    /// MMIO touch and stop chaining there (preserving the interpreter's
+    /// fine-grained MMIO ordering).
+    pub mmio_count: *const u64,
 }
 
 /// Trait for all guest memory access. Every load/store goes through this,

crates/xenia-memory/src/heap.rs

@@ -519,6 +519,7 @@ impl MemoryAccess for GuestMemory {
             page_table: self.page_table.as_ptr() as *const u64,
             mmio_mask: self.mmio_aperture_mask,
             mmio_value: self.mmio_aperture_value,
+            mmio_count: &self.mmio_access_count as *const AtomicU64 as *const u64,
         })
     }