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

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>

hosts/ocaml/bin/sx_server.ml

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

hosts/ocaml/lib/sx_types.ml

@@ -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 *)
 }
 
 

hosts/ocaml/lib/sx_vm.ml

@@ -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
+        (* Walk the closure env chain for inner functions *)
-        with _ -> raise (Eval_error ("VM undefined: " ^ name))
+        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);

shared/sx/ocaml_bridge.py

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