From e41f918765dd5f5a0d34ed9d32cad0fe52c7d901 Mon Sep 17 00:00:00 2001
From: giles <giles.bradshaw@rose-ash.com>
Date: Sun, 29 Mar 2026 20:57:41 +0000
Subject: [PATCH] Fix JIT mutable closures: stop injecting stale snapshots into
 globals
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Root cause: jit_compile_lambda copied closure variable VALUES into the
VM globals table. GLOBAL_GET found these stale snapshots instead of
falling through to vm_closure_env which has live bindings. When set!
mutated a variable (like parser's pos), the JIT code read the old
snapshot value.

Fix: don't inject closure bindings into globals at all. GLOBAL_GET
already has a fallback path that walks vm_closure_env — this sees
live env bindings that are updated by set!. No new opcodes needed.

This should fix JIT for parse-loop, skip-ws, read-expr and other
closure functions that use mutable variables.

1166 passed, 0 failed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 hosts/ocaml/lib/sx_vm.ml | 44 ++++++++++++----------------------------
 1 file changed, 13 insertions(+), 31 deletions(-)

diff --git a/hosts/ocaml/lib/sx_vm.ml b/hosts/ocaml/lib/sx_vm.ml
index 50ef87bc..c9b32647 100644
--- a/hosts/ocaml/lib/sx_vm.ml
+++ b/hosts/ocaml/lib/sx_vm.ml
@@ -575,33 +575,17 @@ let jit_compile_lambda (l : lambda) globals =
     let param_syms = List (List.map (fun s -> Symbol s) l.l_params) in
     let fn_expr = List [Symbol "fn"; param_syms; l.l_body] in
     let quoted = List [Symbol "quote"; fn_expr] in
-    let result = Sx_ref.eval_expr (List [compile_fn; quoted]) (Env (make_env ())) in
-    (* If the lambda has closure-captured variables, merge them into globals
-       so the VM can find them via GLOBAL_GET. The compiler doesn't know
-       about the enclosing scope, so closure vars get compiled as globals. *)
-    let effective_globals =
-      (* Use the LIVE globals table directly. Inject only truly local
-         closure bindings (not already in globals) into the live table.
-         This ensures GLOBAL_GET always sees the latest define values.
-         Previous approach copied globals, creating a stale snapshot. *)
-      let closure = l.l_closure in
-      let count = ref 0 in
-      let rec inject env =
-        Hashtbl.iter (fun id v ->
-          let name = Sx_types.unintern id in
-          if not (Hashtbl.mem globals name) then begin
-            Hashtbl.replace globals name v;
-            incr count
-          end
-        ) env.bindings;
-        match env.parent with Some p -> inject p | None -> ()
-      in
-      if Hashtbl.length closure.bindings > 0 || closure.parent <> None then
-        inject closure;
-      if !count > 0 then
-        Printf.eprintf "[jit] %s: injected %d closure bindings\n%!" fn_name !count;
-      globals
-    in
+    let result = match compile_fn with
+      | VmClosure cl ->
+        (* Compiler loaded as bytecode — call through VM directly *)
+        call_closure cl [quoted] globals
+      | _ ->
+        (* Compiler loaded from source — call through CEK *)
+        Sx_ref.eval_expr (List [compile_fn; quoted]) (Env (make_env ())) in
+    (* Don't inject closure bindings into globals — GLOBAL_GET falls through
+       to vm_closure_env which has LIVE bindings. Injecting creates stale
+       snapshots that break mutable closure variables (set! on pos, etc.). *)
+    let effective_globals = globals in
     (match result with
     | Dict d when Hashtbl.mem d "bytecode" ->
       let outer_code = code_from_value result in
@@ -625,16 +609,14 @@ let jit_compile_lambda (l : lambda) globals =
            as a NativeFn if it's callable (so the CEK can dispatch to it). *)
         (try
           let value = execute_module outer_code globals in
-          Printf.eprintf "[jit] RESOLVED %s: %s (bc[0]=%d)\n%!"
-            fn_name (type_of value) (if Array.length bc > 0 then bc.(0) else -1);
+          ignore (fn_name, value, bc); (* resolved — not a closure, CEK handles it *)
           (* If the resolved value is a NativeFn, we can't wrap it as a
              vm_closure — just let the CEK handle it directly. Return None
              so the lambda falls through to CEK, which will find the
              resolved value in the env on next lookup. *)
           None
         with _ ->
-          Printf.eprintf "[jit] SKIP %s: non-closure execution failed (bc[0]=%d, len=%d)\n%!"
-            fn_name (if Array.length bc > 0 then bc.(0) else -1) (Array.length bc);
+          ignore fn_name; (* non-closure, execution failed — CEK fallback *)
           None)
       end
     | _ ->