diff --git a/hosts/ocaml/bin/sx_server.ml b/hosts/ocaml/bin/sx_server.ml index 00cf1b7..329ffb1 100644 --- a/hosts/ocaml/bin/sx_server.ml +++ b/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 | _ -> diff --git a/hosts/ocaml/lib/sx_types.ml b/hosts/ocaml/lib/sx_types.ml index 3801ee0..5700b45 100644 --- a/hosts/ocaml/lib/sx_types.ml +++ b/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 *) } diff --git a/hosts/ocaml/lib/sx_vm.ml b/hosts/ocaml/lib/sx_vm.ml index 75939bc..e1278df 100644 --- a/hosts/ocaml/lib/sx_vm.ml +++ b/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 - 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); diff --git a/shared/sx/ocaml_bridge.py b/shared/sx/ocaml_bridge.py index fb3d302..184eb3e 100644 --- a/shared/sx/ocaml_bridge.py +++ b/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: