coop/rose-ash
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix JIT mutable closures: stop injecting stale snapshots into globals

Browse Source 
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>
This commit is contained in:
giles
2026-03-29 20:57:41 +00:00
parent 74bab85139
commit e41f918765
1 changed files with 13 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
hosts/ocaml/lib/sx_vm.ml
View File

@@ -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
| _ ->