WIP: fjit: translate writable-and-executable code with a byte check on entry

Builds; not tested yet. Instead of leaving a JIT's code to tier-1, the
function JIT records the bytes every block of such a function covers and
compares them on entry; on a difference it retires their pages
(IC_INVALIDATE) and leaves for resolve, which translates again.
GEX_FJIT_WX=refuse keeps the refusal. Next: smc-direct-jump.c and
smc-threads.c under --fjit (a change made while the function runs, as in
smc-direct-jump, is not seen until its next entry, so that case may need
direct transfers into writable code to leave through dispatch), a
mid-function patch test, then Muck's menu fps with and without.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

crates/gex-fjit/src/emit.rs

@@ -56,7 +56,7 @@ pub enum Linkage<'a> {
     /// anything outside the module is a probe of the cache and a call
     /// through table `table`; a function checks, on entry, that each page
     /// it covers still holds the `PAGE_GEN` value in `pages` (by entry).
-    Runtime { table: u32, pages: &'a HashMap<u64, Vec<(u64, u32)>> },
+    Runtime { table: u32, pages: &'a HashMap<u64, Vec<(u64, u32)>>, guards: &'a crate::Guards },
 }
 
 /// What a function's exits did, for `GEX_FJIT_STATS`.
@@ -206,7 +206,7 @@ pub fn emit(m: &ModuleCx<'_>, f: &Function, loops: &[Loop], stats: &mut Stats) -
     // Its pages unchanged since translation, or the table's newer function
     // for this address (merge-plan §9.8.1; `gex.lookup` retires this one and
     // finds or makes it), before anything else.
-    if let Linkage::Runtime { table, pages } = &m.linkage {
+    if let Linkage::Runtime { table, pages, guards } = &m.linkage {
         let none: &[(u64, u32)] = &[];
         let pages = if cf::exp("NOPAGECHECK") { none } else { pages.get(&f.entry).map(|v| v.as_slice()).unwrap_or(none) };
         for &(page, gen) in pages {
@@ -232,6 +232,61 @@ pub fn emit(m: &ModuleCx<'_>, f: &Function, loops: &[Loop], stats: &mut Stats) -
             fb.return_call_indirect(m.body_ty, *table);
             fb.end();
         }
+        // Code in writable and executable memory: its bytes unchanged
+        // (`crate::Guards`).
+        if let Some(ranges) = guards.get(&f.entry).filter(|r| !r.is_empty()) {
+            let acc = b.local(ValType::I64);
+            let fb = b.f();
+            fb.i64_const(0);
+            fb.local_set(acc);
+            for (start, bytes) in ranges {
+                for (i, chunk) in bytes.chunks(8).enumerate() {
+                    let mut buf = [0u8; 8];
+                    buf[..chunk.len()].copy_from_slice(chunk);
+                    let want = u64::from_le_bytes(buf);
+                    fb.addr_const(start + 8 * i as u64);
+                    fb.load(0x29, 0, 0); // i64.load, unaligned
+                    if chunk.len() < 8 {
+                        fb.i64_const(((1u64 << (8 * chunk.len())) - 1) as i64);
+                        fb.byte(0x83); // i64.and
+                    }
+                    fb.i64_const(want as i64);
+                    fb.byte(0x85); // i64.xor
+                    fb.local_get(acc);
+                    fb.byte(0x84); // i64.or
+                    fb.local_set(acc);
+                }
+            }
+            fb.local_get(acc);
+            fb.byte(0x50); // i64.eqz
+            fb.byte(0x45); // i32.eqz: changed?
+            fb.if_void();
+            let mut pages: Vec<u64> = ranges
+                .iter()
+                .flat_map(|(s, b)| (s / gex_abi::DISPATCH_PAGE)..=((s + b.len().max(1) as u64 - 1) / gex_abi::DISPATCH_PAGE))
+                .collect();
+            pages.sort_unstable();
+            pages.dedup();
+            for p in pages {
+                fb.local_get(0);
+                fb.i64_const(f.entry as i64);
+                fb.i32_const(gex_abi::cache_op::IC_INVALIDATE);
+                fb.i64_const((p * gex_abi::DISPATCH_PAGE) as i64);
+                fb.call(m.cx.cache_op);
+            }
+            fb.local_get(0);
+            fb.addr_i32();
+            fb.i64_const(f.entry as i64);
+            fb.store(0x37, 3, ctx::LAST_TARGET as u32);
+            fb.local_get(0);
+            fb.local_get(1);
+            fb.local_get(2);
+            fb.local_get(3);
+            fb.i32_const(gex_abi::ftab::RESOLVE as i32);
+            fb.ftab_index();
+            fb.return_call_indirect(m.body_ty, *table);
+            fb.end();
+        }
     }
     // The recursion-depth guard, before anything is loaded (§7.1).
     if !m.trapfree.contains_key(&f.entry) {

crates/gex-fjit/src/lib.rs

@@ -105,12 +105,26 @@ pub fn build(
     Ok((wasm, stats))
 }
 
+/// By function entry: `(address, bytes)` ranges of guest code the function
+/// was translated from. Code in writable and executable memory (a JIT's
+/// code area) can change with nothing to announce it -- no page is retired
+/// -- so such a function compares these bytes on entry, all of them (a
+/// method patched in its middle as well as at its head), and on a
+/// difference retires their pages (`cache_op` IC_INVALIDATE, as tier-1's
+/// block guard does) and leaves for the table's resolve, which translates
+/// it again. A change made while the function runs is not seen until its
+/// next entry.
+pub type Guards = HashMap<u64, Vec<(u64, Vec<u8>)>>;
+
 /// A run-time translation of `functions` (memory form, verified; the first
 /// is the one asked for) for the function table: the module, the entries it
 /// holds, what the emitter did. `page_gen` gives the `PAGE_GEN` value of a
 /// page, which each function checks on entry for every page it covers.
 /// `intrinsics` are tier-1's blocks for the entries it serves with a host
 /// builtin (`gex_xlate::intrinsic`): those functions are formed from them.
+/// `guards` are, by entry, the code bytes a function in writable and
+/// executable memory (a JIT's) was translated from, which it compares on
+/// entry (see [`Guards`]).
 pub fn runtime(
     functions: Vec<Function>,
     summaries: &HashMap<u64, Summary>,
@@ -118,6 +132,7 @@ pub fn runtime(
     page_gen: &dyn Fn(u64) -> u32,
     intrinsics: &[Block],
     ret_pops: &HashMap<u64, u64>,
+    guards: &Guards,
 ) -> Result<(Vec<u8>, Vec<u64>, Stats)> {
     let ib: HashMap<u64, &Block> = intrinsics.iter().map(|b| (b.start, b)).collect();
     let mut functions = gex_cgen::select(functions, &ib, &HashSet::new())?;
@@ -139,6 +154,6 @@ pub fn runtime(
         prep::prepare(f, summaries);
     }
     let mut stats = Stats::default();
-    let (wasm, entries) = module::build_runtime(&functions, summaries, class32, pages, ret_pops, &mut stats)?;
+    let (wasm, entries) = module::build_runtime(&functions, summaries, class32, pages, guards, ret_pops, &mut stats)?;
     Ok((wasm, entries, stats))
 }

crates/gex-fjit/src/module.rs

@@ -368,6 +368,7 @@ pub fn build_runtime(
     sums: &HashMap<u64, Summary>,
     class32: bool,
     pages: &HashMap<u64, Vec<(u64, u32)>>,
+    guards: &crate::Guards,
     ret_pops: &HashMap<u64, u64>,
     stats: &mut Stats,
 ) -> Result<(Vec<u8>, Vec<u64>)> {
@@ -426,7 +427,7 @@ pub fn build_runtime(
             trapfree: &trapfree,
             ret_sp_delta: cf::ret_sp_delta(class32),
             ret_pops,
-            linkage: crate::emit::Linkage::Runtime { table: 0, pages },
+            linkage: crate::emit::Linkage::Runtime { table: 0, pages, guards },
             body_ty,
         };
         let mut emitted: Vec<(u64, Emitted)> = Vec::new();

crates/gex-fjit/tests/replay.rs

@@ -27,6 +27,6 @@ fn replay() {
         }
     }
     let sums = std::collections::HashMap::new();
-    let r = gex_fjit::runtime(funcs, &sums, true, &|_| 1, &[], &std::collections::HashMap::new());
+    let r = gex_fjit::runtime(funcs, &sums, true, &|_| 1, &[], &std::collections::HashMap::new(), &Default::default());
     println!("{:?}", r.map(|(w, e, _)| (w.len(), e.len())));
 }

crates/gex-xlate/src/fjit.rs

@@ -131,6 +131,7 @@ impl Jit {
         let code = Code::memory(mem, executable, &ro);
         let opts = RecoverOptions::default();
         let mut batch: BTreeMap<u64, Function> = BTreeMap::new();
+        let mut guards = gex_fjit::Guards::new();
         let mut order: Vec<u64> = Vec::new();
         let mut work: VecDeque<u64> = VecDeque::from([addr]);
         let limit = batch_limit();
@@ -159,13 +160,30 @@ impl Jit {
             // direct jump reached when it was translated -- in a JIT's
             // fresh area, often bytes not yet written. It stays tier-1.
             if f.blocks.iter().any(|b| !b.synthetic && (volatile(b.start) || volatile(b.end.max(b.start + 1) - 1))) {
-                if std::env::var_os("GEX_FJIT_STATS").is_some() {
-                    eprintln!("[fjit] {a:#x}: code in writable memory; left to the fallback");
+                if wx_refused() {
+                    if std::env::var_os("GEX_FJIT_STATS").is_some() {
+                        eprintln!("[fjit] {a:#x}: code in writable memory; left to the fallback");
+                    }
+                    if a == addr {
+                        return None;
+                    }
+                    continue;
                 }
-                if a == addr {
-                    return None;
+                // Its bytes, compared on entry (`gex_fjit::Guards`).
+                let mut spans: Vec<(u64, u64)> = f.blocks.iter().filter(|b| !b.synthetic).map(|b| (b.start, b.end.max(b.start + 1))).collect();
+                spans.sort_unstable();
+                let mut merged: Vec<(u64, u64)> = Vec::new();
+                for (lo, hi) in spans {
+                    match merged.last_mut() {
+                        Some(last) if lo <= last.1 => last.1 = last.1.max(hi),
+                        _ => merged.push((lo, hi)),
+                    }
                 }
-                continue;
+                let ranges: Vec<(u64, Vec<u8>)> = merged
+                    .into_iter()
+                    .filter_map(|(lo, hi)| mem.get(lo as usize..hi as usize).map(|b| (lo, b.to_vec())))
+                    .collect();
+                guards.insert(a, ranges);
             }
             if !intrinsics.contains_key(&a) {
                 for c in f.call_sites() {
@@ -228,7 +246,7 @@ impl Jit {
         // A bug in the emitter costs the batch its translation, not the
         // process: the fallback runs it.
         let emitted = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
-            gex_fjit::runtime(functions, &sums, class32, page_gen, &builtins, &pops)
+            gex_fjit::runtime(functions, &sums, class32, page_gen, &builtins, &pops, &guards)
         }))
         .map_err(|p| p.downcast_ref::<String>().cloned().or_else(|| p.downcast_ref::<&str>().map(|s| s.to_string())).unwrap_or_default());
         let out = match emitted {
@@ -357,3 +375,11 @@ impl State {
     }
 }
 
+
+/// `GEX_FJIT_WX=refuse`: leave code in writable and executable memory to
+/// tier-1's block guards instead of translating it with a byte check on
+/// entry (`gex_fjit::Guards`).
+fn wx_refused() -> bool {
+    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+    *ON.get_or_init(|| std::env::var("GEX_FJIT_WX").as_deref() == Ok("refuse"))
+}