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

VM: closure env chain for GLOBAL_GET/SET + remove JIT skip

Browse Source 
vm_closure now stores the original closure env (vm_closure_env).
GLOBAL_GET walks the closure env chain when the variable isn't in
vm.globals. GLOBAL_SET writes to the correct env in the chain.

This enables JIT compilation of all named functions regardless of
closure depth. No more closure skip check needed.

Pre-compile time back to ~7s (was 37s with closure skip).

Note: sx-parse sibling list parsing still has issues — the root
cause is in how the JIT-compiled letrec + OP_CLOSURE interacts
with the upvalue cell mechanism. Investigation ongoing.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
giles
2026-03-23 21:23:27 +00:00
parent 30cfbf777a
commit 8958714c85
4 changed files with 30 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
hosts/ocaml/bin/sx_server.ml
View File

@@ -913,9 +913,6 @@ let register_jit_hook env =
| Some _ -> None (* compile failed — CEK handles *)
| None ->
if !_jit_compiling then None
else if l.l_closure.bindings != env.bindings then
(* Skip JIT for inner functions — closure != globals *)
None
else begin
let fn_name = match l.l_name with Some n -> n | None -> "?" in
_jit_compiling := true;
@@ -1300,7 +1297,7 @@ let rec dispatch env cmd =
try
Sx_vm.call_closure
{ vm_code = code; vm_upvalues = [||]; vm_name = lam.l_name;
vm_env_ref = live_env }
vm_env_ref = live_env; vm_closure_env = None }
args live_env
with
| _ ->

1
hosts/ocaml/lib/sx_types.ml
View File

@@ -134,6 +134,7 @@ and vm_closure = {
vm_upvalues : vm_upvalue_cell array;
vm_name : string option;
vm_env_ref : (string, value) Hashtbl.t;
vm_closure_env : env option; (** Original closure env for inner functions *)
}

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

@@ -36,7 +36,7 @@ let jit_compile_ref : (lambda -> (string, value) Hashtbl.t -> vm_closure option)
Prevents retrying compilation on every call. *)
let jit_failed_sentinel = {
vm_code = { vc_arity = -1; vc_locals = 0; vc_bytecode = [||]; vc_constants = [||] };
vm_upvalues = [||]; vm_name = Some "__jit_failed__"; vm_env_ref = Hashtbl.create 0
vm_upvalues = [||]; vm_name = Some "__jit_failed__"; vm_env_ref = Hashtbl.create 0; vm_closure_env = None
}
let is_jit_failed cl = cl.vm_code.vc_arity = -1
@@ -154,14 +154,35 @@ let rec run vm =
let idx = read_u16 frame in
let name = match consts.(idx) with String s -> s | _ -> "" in
let v = try Hashtbl.find vm.globals name with Not_found ->
try Sx_primitives.get_primitive name
with _ -> raise (Eval_error ("VM undefined: " ^ name))
(* Walk the closure env chain for inner functions *)
let rec env_lookup e =
try Hashtbl.find e.bindings name
with Not_found ->
match e.parent with Some p -> env_lookup p | None ->
try Sx_primitives.get_primitive name
with _ -> raise (Eval_error ("VM undefined: " ^ name))
in
match frame.closure.vm_closure_env with
| Some env -> env_lookup env
| None ->
try Sx_primitives.get_primitive name
with _ -> raise (Eval_error ("VM undefined: " ^ name))
in
push vm v; run vm
| 21 (* OP_GLOBAL_SET *) ->
let idx = read_u16 frame in
let name = match consts.(idx) with String s -> s | _ -> "" in
Hashtbl.replace vm.globals name (peek vm);
(* Write to closure env if the name exists there (mutable closure vars) *)
let written = match frame.closure.vm_closure_env with
| Some env ->
let rec find_env e =
if Hashtbl.mem e.bindings name then
(Hashtbl.replace e.bindings name (peek vm); true)
else match e.parent with Some p -> find_env p | None -> false
in find_env env
| None -> false
in
if not written then Hashtbl.replace vm.globals name (peek vm);
run vm
(* ---- Control flow ---- *)
@@ -235,7 +256,7 @@ let rec run vm =
(* Capture from enclosing frame's upvalue — already a shared cell *)
frame.closure.vm_upvalues.(index)
) in
let cl = { vm_code = code; vm_upvalues = upvalues; vm_name = None; vm_env_ref = vm.globals } in
let cl = { vm_code = code; vm_upvalues = upvalues; vm_name = None; vm_env_ref = vm.globals; vm_closure_env = None } in
(* Wrap as NativeFn that calls back into the VM *)
let fn = NativeFn ("vm-closure", fun args ->
call_closure cl args vm.globals)
@@ -328,10 +349,6 @@ and vm_call vm f args =
push vm (Sx_ref.cek_call f (List args))
| None ->
if l.l_name <> None
(* Skip JIT for inner functions (closure != globals).
The closure merging produces incorrect variable resolution
for functions that capture letrec/let-local bindings. *)
&& l.l_closure.bindings == vm.globals
then begin
(* Pre-mark before compile attempt to prevent re-entrancy *)
l.l_compiled <- Some jit_failed_sentinel;
@@ -391,7 +408,7 @@ and call_closure cl args globals =
(** Execute a compiled module (top-level bytecode). *)
let execute_module code globals =
let cl = { vm_code = code; vm_upvalues = [||]; vm_name = Some "module"; vm_env_ref = globals } in
let cl = { vm_code = code; vm_upvalues = [||]; vm_name = Some "module"; vm_env_ref = globals; vm_closure_env = None } in
let vm = create globals in
let frame = { closure = cl; ip = 0; base = 0; local_cells = Hashtbl.create 4 } in
for _ = 0 to code.vc_locals - 1 do push vm Nil done;
@@ -454,7 +471,7 @@ let jit_compile_lambda (l : lambda) globals =
let inner_val = outer_code.vc_constants.(idx) in
let code = code_from_value inner_val in
Some { vm_code = code; vm_upvalues = [||];
vm_name = l.l_name; vm_env_ref = effective_globals }
vm_name = l.l_name; vm_env_ref = effective_globals; vm_closure_env = Some l.l_closure }
else begin
Printf.eprintf "[jit] FAIL %s: closure index %d out of bounds (pool=%d)\n%!"
fn_name idx (Array.length outer_code.vc_constants);

10
shared/sx/ocaml_bridge.py
View File

@@ -440,16 +440,6 @@ class OcamlBridge:
_logger.warning("OCaml load skipped %s: %s",
filepath, e)
# sx-parse has deeply nested letrec + define inside fn
# bodies. The JIT closure/upvalue mechanism can't handle
# the multiple nesting layers correctly — nested list
# parsing produces wrong results. Force CEK.
try:
await self._send('(vm-reset-fn "sx-parse")')
await self._read_until_ok(ctx=None)
except OcamlBridgeError:
pass
# SSR overrides: effect is a no-op on the server (prevents
# reactive loops during island SSR — effects are DOM side-effects)
try: