fjit: forget learned state over retired code; misses do not make starts; fjit.rs out of the unit key

Two holes in the function JIT's learned state (gex_xlate::fjit), both
found in the page with Hollow Knight (docs/unity-dotnet.md item 34) and
both present natively:

- What recovery learned (function starts, never-returning functions,
  functions translated, read-only ranges) is keyed by address and
  outlived the code: a dlclose and a later dlopen put another library at
  the same addresses. HK's SDL loads libdbus, unloads it, and maps libxcb
  over it; a libdbus start ended libxcb's _xcb_in_read_block before its
  cmovle as a tail call with the flags dead, so min(queue_len, len) read
  stale flags (the X setup reply read from an empty queue: SDL's video
  init failed, a black screen; with the unit, a memmove of -8 bytes).
  Jit::retire(process, ranges) drops it; the Translator trait gains
  retire(), which the native host's ftab::on_retire and the page's
  retire_unmapped_code call (the page's epoch workaround goes).
- Every translation root was recorded as a strong start for good, and a
  root can be the middle of a function (a return or a signal's resume
  through dispatch, a computed goto's label, a guessed slot target),
  splitting the function around it the same way. A root stays a start
  only if something (a call, the image's symbols) made it one already.

With misses no longer making starts, recovery can walk over a JIT
function's entry inside the function it recovers; the page asked its
"served by the unit or the JIT" test of instruction fetch too (as part
of `executable`), read those bytes as non-code and decoded them as
illegal (HK without the unit: SIGILL at start-up, hb52). The test now
bounds the batch only: Jit::translate_skipping takes it as `skip`
(translate_in passes none).

tests/fjit-remap: code mapped, run, unmapped and replaced at the same
address. gexc built before this gets 99,800 of 200,000 min() results
wrong under --fjit; with it, output matches native.

gex-cgen/build.rs leaves fjit.rs out of GEX_CGEN_RECOVERY_INPUTS: the
function JIT runs recovery at run time and no unit calls it. This
re-keys every unit cache once; after it a change to fjit.rs leaves the
key alone (checked: c6014ffc... before and after an edit to fjit.rs,
855acf6a... after one to func.rs).

In the page (edger), HK with full fjit and no unit (hb53): the menu,
"window backend is x11", first frame 178 s after the page opened and
35.7 presents/s at 392 s, against the epoch shim's hb51: 222 s and 33.9.

Checked on gooner: tests/fjit-remap, tests/fjit-procs, tests/x86-programs
(82 ok), tests/unity-dotnet/embed (PASS), gex-xlate's tests.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_1896c5b302a54696aaa92c5f4047f20b

crates/gex-cgen/build.rs

@@ -3,8 +3,10 @@
 //!
 //! * `GEX_CGEN_RECOVERY_INPUTS`: what decides an image's recovery (its GIR
 //!   and summaries) -- the IR, the lifters and decoders, gex-xlate's recovery
-//!   (everything in it but the runtime translator's `lib.rs` and the
-//!   profile and validation code), gex-abi, and gex-ld (placing the image);
+//!   (everything in it but the runtime translator's `lib.rs` and
+//!   `fjit.rs` -- the function JIT drives recovery at run time, a unit
+//!   never calls it -- and the profile and validation code), gex-abi, and
+//!   gex-ld (placing the image);
 //! * `GEX_CGEN_UNIT_INPUTS`: what decides the C and the module from a
 //!   recovery -- this crate (and its runtime). What the C takes from gex-abi
 //!   is hashed at run time (`image::abi_fingerprint`: the header and every
@@ -55,7 +57,7 @@ fn main() {
         rec.extend(src(c));
     }
     rec.extend(src("gex-xlate").into_iter().filter(|p| {
-        !matches!(p.file_name().and_then(|n| n.to_str()), Some("lib.rs" | "profile.rs" | "validate.rs"))
+        !matches!(p.file_name().and_then(|n| n.to_str()), Some("lib.rs" | "profile.rs" | "validate.rs" | "fjit.rs"))
     }));
     // gex-abi is not hashed here: what the C takes from it is the generated
     // header and the values `image::abi_fingerprint` lists, which the unit

crates/gex-kernel/src/ftab.rs

@@ -258,6 +258,10 @@ pub trait Translator: Send + Sync {
     fn image(&self, _process: u64, _program: &gex_ld::Program, _bytes: Option<&[u8]>) {}
     /// Guest process `process` is gone (its last reference dropped).
     fn forget(&self, _process: u64) {}
+    /// Guest process `process` unmapped or reprotected the code in
+    /// `ranges` (`Kernel::retired_code`): what the translator learned of
+    /// those addresses no longer describes them.
+    fn retire(&self, _process: u64, _ranges: &[(u64, u64)]) {}
 }
 
 static TRANSLATOR: std::sync::OnceLock<Box<dyn Translator>> = std::sync::OnceLock::new();

crates/gex-xlate/src/fjit.rs

@@ -87,6 +87,29 @@ impl Jit {
         self.add_image_in(process, image, elf);
     }
 
+    /// Process `process` unmapped or reprotected the code in `ranges`:
+    /// what was learned there (function starts, never-returning functions,
+    /// functions translated, read-only ranges) described the code that was
+    /// there, and a `dlclose` and a later `dlopen` put another library at
+    /// the same addresses. Kept, the old library's starts split the new
+    /// one's functions where it had none: HK's SDL loads libdbus, unloads
+    /// it, and maps libxcb over it; a libdbus start ended libxcb's
+    /// `_xcb_in_read_block` before its `cmovle` as a tail call, flags dead,
+    /// and it read the X setup reply out of an empty queue.
+    pub fn retire(&self, process: u64, ranges: &[(u64, u64)]) {
+        let mut st = self.st.lock().unwrap_or_else(|e| e.into_inner());
+        let Some(l) = st.procs.get_mut(&process) else { return };
+        for &(lo, hi) in ranges.iter().filter(|r| r.1 > r.0) {
+            let gone: Vec<u64> = l.starts.range(lo..hi).map(|(&a, _)| a).collect();
+            for a in gone {
+                l.starts.remove(&a);
+            }
+            l.noreturn.retain(|&a| a < lo || a >= hi);
+            l.translated.retain(|&a| a < lo || a >= hi);
+            l.readonly.retain(|&(s, e)| e <= lo || s >= hi);
+        }
+    }
+
     /// Process `process` is gone.
     pub fn forget(&self, process: u64) {
         let mut st = self.st.lock().unwrap_or_else(|e| e.into_inner());
@@ -159,6 +182,28 @@ impl Jit {
         volatile: &dyn Fn(u64) -> bool,
         page_gen: &dyn Fn(u64) -> u32,
         class32: bool,
+    ) -> Option<Vec<u8>> {
+        self.translate_skipping(process, mem, addr, executable, &|_| false, read_only, volatile, page_gen, class32)
+    }
+
+    /// [`Jit::translate_in`], leaving out of the batch the callees `skip`
+    /// names (served by another tier: the call reaches them through the
+    /// table). `skip` bounds the batch only, never instruction fetch: a
+    /// function another tier serves can have an entry in the middle of the
+    /// one being recovered, and fetch taking its bytes for non-code decodes
+    /// them as illegal.
+    #[allow(clippy::too_many_arguments)]
+    pub fn translate_skipping(
+        &self,
+        process: u64,
+        mem: &[u8],
+        addr: u64,
+        executable: &dyn Fn(u64) -> bool,
+        skip: &dyn Fn(u64) -> bool,
+        read_only: &dyn Fn(u64, u64) -> bool,
+        volatile: &dyn Fn(u64) -> bool,
+        page_gen: &dyn Fn(u64) -> u32,
+        class32: bool,
     ) -> Option<Vec<u8>> {
         if gex_abi::PPC || (class32 && !gex_abi::X86) {
             return None;
@@ -188,11 +233,31 @@ impl Jit {
             if batch.len() >= limit {
                 break;
             }
-            if batch.contains_key(&a) || (a != addr && st.translated.contains(&a)) || !executable(a) {
+            if batch.contains_key(&a) || (a != addr && (st.translated.contains(&a) || skip(a))) || !executable(a) {
                 continue;
             }
             let Learned { starts, noreturn, .. } = &mut *st;
+            // Recovery records `a` as a function start for good, which bounds
+            // every function recovered after it: a later recovery of the
+            // function around `a` ends there as at a tail call, with the
+            // flags dead. A root the table missed at can be the middle of a
+            // function (a return or a signal's resume through dispatch, a
+            // computed goto's label, a guessed slot target), so `a` stays a
+            // start only if something already made it one (a call, the
+            // image's symbols). The page's JIT, with no images to seed it,
+            // split libxcb's `_xcb_in_read_block` before a `cmovle` this way.
+            let before = starts.get(&a).cloned();
             let f = func::recover_one(&code, a, starts, noreturn, &opts);
+            if !before.as_ref().is_some_and(|b| b.strong()) {
+                match before {
+                    Some(b) => {
+                        starts.insert(a, b);
+                    }
+                    None => {
+                        starts.remove(&a);
+                    }
+                }
+            }
             if let Err(e) = f.verify().and_then(|_| block_local(&f)) {
                 if std::env::var_os("GEX_FJIT_STATS").is_some() {
                     eprintln!("[fjit] {a:#x}: recovered malformed ({e}); left to the fallback");

crates/gexc/src/main.rs

@@ -560,6 +560,10 @@ impl gex_kernel::ftab::Translator for FjitTranslator {
     fn forget(&self, process: u64) {
         self.0.forget(process);
     }
+
+    fn retire(&self, process: u64, ranges: &[(u64, u64)]) {
+        self.0.retire(process, ranges);
+    }
 }
 
 /// The guest memory for this run (docs/mem64.md §5): `--mem64`/`--mem32`,

shims/browser/src/fjit.rs

@@ -16,7 +16,7 @@
 //! run on worker threads (`GEX_FJIT_BACKGROUND`); here it is always inline.
 
 use gex_kernel::GuestMem;
-use std::collections::{BTreeSet, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Mutex, OnceLock};
 
@@ -326,48 +326,17 @@ fn mid_function(addr: u64) -> bool {
     mid
 }
 
-/// What the JIT learns (function starts, functions translated, functions
-/// that never return) is kept under this key, which moves on whenever code
-/// is unmapped: it describes addresses, and a `dlclose` and a later
-/// `dlopen` put another library at them. HK's SDL loads libdbus, unloads it,
-/// and maps libxcb over the same range; with libdbus's starts kept, libxcb's
-/// `_xcb_in_read_block` ended at the address of a libdbus function (+0x30, the
-/// `cmovle` of `min(queue_len, len)`) as at a tail call, flags dead, and
-/// read the setup reply out of an empty queue: X11 video init failed, the
-/// black screen (hb26, hb45; hs5 `scanconst=` showed the start was
-/// libdbus's), and with the unit, a memmove of -8 bytes (hp8).
-static EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
-
-/// The 4 KiB pages holding a root or a function translated this epoch: the
-/// starts the JIT learned lie in them (a batch's callees in other objects
-/// are its members too), so code unmapped elsewhere leaves them be. ALSA
-/// loads and unloads its plugins over and over at start-up (six of hb50's
-/// ten epochs), and each forgetting cost a re-translation of everything
-/// called after it (libasound: 27 MiB of wasm by HK's menu in hp11).
-static TOUCHED: Mutex<Option<BTreeSet<u64>>> = Mutex::new(None);
-
-fn touch(a: u64) {
-    lock(&TOUCHED).get_or_insert_with(Default::default).insert(a >> 12);
-}
-
 /// Code in `ranges` was unmapped or reprotected (`retire_unmapped_code`):
-/// start learning afresh, if this epoch learned anything there.
+/// what the JIT learned there goes (`Jit::retire`). HK's SDL loads libdbus,
+/// unloads it, and maps libxcb over the same range; with libdbus's starts
+/// kept, libxcb's `_xcb_in_read_block` ended at a libdbus function's address
+/// (+0x30, the `cmovle` of `min(queue_len, len)`) as at a tail call, flags
+/// dead, and read the setup reply out of an empty queue: X11 video init
+/// failed, the black screen (hb26, hb45; `scanconst=` in hs5 showed the
+/// start was libdbus's), and with the unit, a memmove of -8 bytes (hp8).
 pub(crate) fn on_retire(ranges: &[(u64, u64)]) {
-    {
-        let mut t = lock(&TOUCHED);
-        let set = t.get_or_insert_with(Default::default);
-        if !ranges.iter().any(|&(s, e)| e > s && set.range(s >> 12..=(e - 1) >> 12).next().is_some()) {
-            return;
-        }
-        set.clear();
-    }
-    let old = EPOCH.fetch_add(1, Ordering::AcqRel);
     if let Some(jit) = JIT.get() {
-        jit.forget(old);
-        if old < 16 || (old + 1).is_power_of_two() {
-            let r: Vec<String> = ranges.iter().take(4).map(|&(s, e)| format!("{s:#x}-{e:#x}")).collect();
-            crate::warn_line(&format!("[fjit] code unmapped ({}): learning afresh, epoch {}\n", r.join(" "), old + 1));
-        }
+        jit.retire(0, ranges);
     }
 }
 
@@ -408,28 +377,26 @@ fn translate_root(addr: u64) -> Option<u32> {
     // An address an ELF provider holds is not code to translate (cgen::elfprov).
     // With `?fjitlibs=` rules, callees outside them are not brought in either.
     let rules = !lock(&LIBS).is_empty();
-    // A callee a unit already serves is not brought into the batch: the call
-    // reaches it through the table. Without this, a miss in UnityPlayer (a
-    // function its unit lacks) re-translated up to 63 of the unit's own
-    // functions with it -- 3,667 functions and 36 MiB of wasm by HK's menu
-    // (ht3), JIT code shadowing clang's. Only the unit tier: `executable` is
-    // also what instruction fetch asks of every byte, and the JIT publishes
-    // resume points inside functions, so counting its own entries made it
-    // decode `cmovle` in libxcb's `_xcb_in_read_block` as illegal (hu15).
-    // (That entry was a function of libdbus, unloaded since: the epoch now
-    // forgets those. A current JIT function's entry counts again: after a
-    // forgetting the JIT no longer knows what it made, and re-translated
-    // libasound into 28 MiB of wasm by the menu, hp12. Fetch meeting one
-    // inside another function would need an entry in the middle of a
-    // function, which misses no longer make, `mid_function`.)
+    // A callee a unit or the JIT already serves is not brought into the
+    // batch: the call reaches it through the table. Without this, a miss in
+    // UnityPlayer (a function its unit lacks) re-translated up to 63 of the
+    // unit's own functions with it -- 3,667 functions and 36 MiB of wasm by
+    // HK's menu (ht3), JIT code shadowing clang's -- and libasound's served
+    // callees went into every batch (28 MiB of wasm by the menu, hp12).
+    // `skip` bounds the batch only. Asked of instruction fetch (as part of
+    // `executable`), an entry inside the function being recovered -- a
+    // resume point, or a root the JIT took inside a function and, since
+    // misses no longer make starts, does not stop recovery at -- read as
+    // non-code and decoded as illegal: `cmovle` in libxcb's
+    // `_xcb_in_read_block` (hu15), and SIGILL at HK's start-up without the
+    // unit (hb52).
     let served_memo = std::cell::RefCell::new(std::collections::HashMap::<u64, bool>::new());
-    let served = |a: u64| a != addr && *served_memo.borrow_mut().entry(a).or_insert_with(|| crate::cgen::unit_serves(a) || crate::cgen::jit_serves(a));
+    let served = |a: u64| *served_memo.borrow_mut().entry(a).or_insert_with(|| crate::cgen::unit_serves(a) || crate::cgen::jit_serves(a));
     let executable = |a: u64| {
         exec.iter().any(|&(s, e)| a >= s && a < e)
             && !crate::cgen::elfprov::provided(a)
             && !crate::cgen::libprov::replaced(a)
             && (!rules || lib_wanted(a))
-            && !served(a)
     };
     let read_only = |a: u64, n: u64| ro.iter().any(|&(s, e)| a >= s && a.saturating_add(n) <= e);
     // Writable-and-executable code (a JIT's) stays tier-1, as natively.
@@ -446,11 +413,7 @@ fn translate_root(addr: u64) -> Option<u32> {
     };
     let class32 = gex_abi::X86 && gex_xlate::arch::compat();
     trace(&format!("translate {addr:#x}"));
-    let epoch = EPOCH.load(Ordering::Acquire);
-    // What this epoch learns about (the root, even if it fails: recovery
-    // records it as a start first).
-    touch(addr);
-    let Some(wasm) = jit.translate_in(epoch, crate::guest_bytes(), addr, &executable, &read_only, &volatile, &pg, class32) else {
+    let Some(wasm) = jit.translate_skipping(0, crate::guest_bytes(), addr, &executable, &served, &read_only, &volatile, &pg, class32) else {
         DECLINED.fetch_add(1, Ordering::Relaxed);
         trace(&format!("declined {addr:#x}"));
         return None;
@@ -460,9 +423,6 @@ fn translate_root(addr: u64) -> Option<u32> {
         DECLINED.fetch_add(1, Ordering::Relaxed);
         return None;
     }
-    for &f in &funcs {
-        touch(f);
-    }
     note_object(addr, funcs.len(), wasm.len());
     dump_module(addr, &funcs, &wasm);
     scan_const(addr, &funcs, &wasm);

shims/native/src/ftab.rs

@@ -1004,6 +1004,9 @@ pub(crate) fn install_bridges(mut cx: impl AsContextMut<Data = HostState>, slots
 /// sending callers there first.
 pub(crate) fn on_retire(sh: &mut JitShared, memory: &mut [u8], ranges: &[(u64, u64)]) {
     sh.ftab.epoch += 1;
+    if let Some(t) = gex_kernel::ftab::translator() {
+        t.retire(sh.ftab.process.0, ranges);
+    }
     sh.ftab.queued.retain(|&a| !ranges.iter().any(|&(s, e)| a >= s && a < e));
     for &(start, end) in ranges {
         let hit: Vec<(u64, gex_kernel::ftab::Func)> = sh.ftab.map.range(start, end).collect();

tests/fjit-remap/remap.c

@@ -0,0 +1,55 @@
+// The function JIT must forget what it learned of code that is unmapped: a
+// dlclose and a later dlopen put another library at the same addresses. Code
+// A has a function at +0x10; code B, mapped over it, has a min() whose cmp is
+// before +0x10 and whose cmovle is at it. Kept, A's start split B's function
+// there as at a tail call, with the flags dead (HK's libdbus and libxcb in the
+// page, docs/unity-dotnet.md). Output must match native.
+#include <stdio.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <unistd.h>
+
+typedef int (*fn)(int, int);
+
+static void *at;
+
+// The first call picks the address; the second maps over the same one.
+static void *place(const unsigned char *code, size_t n) {
+	void *p = mmap(at, 4096, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | (at ? MAP_FIXED : 0), -1, 0);
+	if (p == MAP_FAILED || (at && p != at)) return 0;
+	at = p;
+	memset(p, 0xcc, 4096);
+	memcpy(p, code, n);
+	if (mprotect(p, 4096, PROT_READ | PROT_EXEC)) return 0;
+	return p;
+}
+
+int main(void) {
+	// A: +0x00 ret; +0x10 lea (%rdi,%rsi),%eax; ret
+	unsigned char a[0x14] = {0xc3};
+	memset(a + 1, 0x90, 0xf);
+	a[0x10] = 0x8d; a[0x11] = 0x04; a[0x12] = 0x37; a[0x13] = 0xc3;
+	char *p = place(a, sizeof a);
+	if (!p) { puts("map A failed"); return 1; }
+	volatile fn f = (fn)(p + 0x10);
+	long s = 0;
+	for (int i = 0; i < 200000; i++) s += f(i, 1);
+	printf("A %ld\n", s);
+	if (munmap(p, 4096)) { puts("munmap failed"); return 1; }
+	// B: mov %esi,%eax; cmp %esi,%edi; nop x12; +0x10 cmovle %edi,%eax; ret
+	unsigned char b[0x14] = {0x89, 0xf0, 0x39, 0xf7};
+	memset(b + 4, 0x90, 0xc);
+	b[0x10] = 0x0f; b[0x11] = 0x4e; b[0x12] = 0xc7; b[0x13] = 0xc3;
+	p = place(b, sizeof b);
+	if (!p) { puts("map B failed"); return 1; }
+	volatile fn g = (fn)p;
+	long bad = 0, t = 0;
+	for (int i = 0; i < 200000; i++) {
+		int x = (i * 7919) % 1000 - 500, y = (i * 104729) % 1000 - 500;
+		int m = g(x, y);
+		t += m;
+		if (m != (x < y ? x : y)) bad++;
+	}
+	printf("B %ld bad %ld\n", t, bad);
+	return 0;
+}

tests/fjit-remap/run.sh

@@ -0,0 +1,22 @@
+#!/bin/sh
+# Code unmapped and replaced at the same address (remap.c): output under gexc
+# --fjit must match native.
+#
+#   tests/fjit-remap/run.sh [gexc]      (needs zig for the musl target)
+set -u
+HERE=$(cd "$(dirname "$0")" && pwd)
+GEXC=${1:-target/x86/release/gexc}
+W=${WORK:-$(mktemp -d)}
+mkdir -p "$W"
+zig cc -target x86_64-linux-musl -O1 -static -o "$W/remap" "$HERE/remap.c" || exit 1
+native=$("$W/remap")
+fail=0
+for flags in "" "--fjit"; do
+	gex=$("$GEXC" run $flags "$W/remap" 2>&1)
+	if [ "$native" = "$gex" ]; then
+		echo "ok   remap ${flags:-t1}"
+	else
+		echo "FAIL remap ${flags:-t1}"; printf 'native:\n%s\ngex:\n%s\n' "$native" "$gex"; fail=1
+	fi
+done
+exit $fail