From 3a9d1135371a8c3bd66d669619921711acea7093 Mon Sep 17 00:00:00 2001 From: giles Date: Mon, 6 Apr 2026 12:08:08 +0000 Subject: [PATCH] Fix JIT mutable closure bug: vm-global-get now checks closure env first MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit vm-global-get checked vm.globals before closure-env, while vm-global-set wrote to closure-env first. This asymmetry meant set! mutations to mutable closure variables (e.g. parser position counters) were invisible to sibling closures reading via JIT — they saw stale snapshots in the globals table. Reversed vm-global-get lookup order: closure env → globals → primitives, matching vm-global-set. Also enabled JIT in the MCP harness (compiler.sx loading, env_bind hook for live globals sync, jit_try_call hook) so sx_harness_eval exercises the same code path as the server. Co-Authored-By: Claude Opus 4.6 (1M context) --- hosts/ocaml/bin/mcp_tree.ml | 52 +++++++++++++++++++++++++++++++++++++ lib/vm.sx | 27 +++++++++++-------- 2 files changed, 68 insertions(+), 11 deletions(-) diff --git a/hosts/ocaml/bin/mcp_tree.ml b/hosts/ocaml/bin/mcp_tree.ml index e52448ea..dfd8f8d2 100644 --- a/hosts/ocaml/bin/mcp_tree.ml +++ b/hosts/ocaml/bin/mcp_tree.ml @@ -27,6 +27,48 @@ let load_sx_file e path = ignore (Sx_ref.eval_expr expr (Env e)) ) exprs +(* JIT infrastructure — shared VM globals table, kept in sync via env_bind hook *) +let _mcp_vm_globals : (string, value) Hashtbl.t = Hashtbl.create 2048 +let _jit_compiling = ref false +let _jit_warned : (string, bool) Hashtbl.t = Hashtbl.create 32 + +let register_mcp_jit_hook () = + Sx_runtime._jit_try_call_fn := Some (fun f args -> + match f with + | Lambda l -> + (match l.l_compiled with + | Some cl when not (Sx_vm.is_jit_failed cl) -> + (try Some (Sx_vm.call_closure cl args cl.vm_env_ref) + with e -> + let fn_name = match l.l_name with Some n -> n | None -> "?" in + if not (Hashtbl.mem _jit_warned fn_name) then begin + Hashtbl.replace _jit_warned fn_name true; + Printf.eprintf "[mcp-jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e) + end; + None) + | Some _ -> None + | None -> + if !_jit_compiling then None + else begin + let fn_name = match l.l_name with Some n -> n | None -> "?" in + if Hashtbl.mem _jit_warned fn_name then None + else begin + _jit_compiling := true; + let compiled = Sx_vm.jit_compile_lambda l _mcp_vm_globals in + _jit_compiling := false; + match compiled with + | Some cl -> + l.l_compiled <- Some cl; + (try Some (Sx_vm.call_closure cl args cl.vm_env_ref) + with e -> + Printf.eprintf "[mcp-jit] %s first-call fallback: %s\n%!" fn_name (Printexc.to_string e); + Hashtbl.replace _jit_warned fn_name true; + None) + | None -> None + end + end) + | _ -> None) + let setup_env () = let e = make_env () in (* Primitives are auto-registered at module init *) @@ -35,6 +77,11 @@ let setup_env () = match v with | Thunk (body, closure_env) -> Sx_ref.eval_expr body (Env closure_env) | other -> other); + (* JIT: mirror root-env bindings into shared VM globals table *) + Sx_types._env_bind_hook := Some (fun env name v -> + if env.parent = None then + if not (Sx_primitives.is_primitive name) then + Hashtbl.replace _mcp_vm_globals name v); (* Character classification for parser *) let bind name fn = ignore (env_bind e name (NativeFn (name, fn))) in bind "is-whitespace?" (fun args -> match args with @@ -296,6 +343,11 @@ let setup_env () = (NativeFn ("html:" ^ tag, fun args -> List (Symbol tag :: args)))) ) Sx_render.html_tags; ignore (Sx_types.env_bind e "island?" (NativeFn ("island?", fun args -> match args with [Island _] -> Bool true | _ -> Bool false))); + (* Load compiler + enable JIT *) + (try load_sx_file e (Filename.concat lib_dir "compiler.sx"); + register_mcp_jit_hook (); + Printf.eprintf "[mcp] JIT enabled (compiler.sx loaded)\n%!" + with exn -> Printf.eprintf "[mcp] Warning: compiler.sx load failed (JIT disabled): %s\n%!" (Printexc.to_string exn)); Printf.eprintf "[mcp] SX tree-tools + harness + eval-rules + render loaded\n%!"; env := e diff --git a/lib/vm.sx b/lib/vm.sx index fbd63557..4f90234e 100644 --- a/lib/vm.sx +++ b/lib/vm.sx @@ -297,27 +297,32 @@ vm-global-get (fn (vm frame name) - "Look up a global: globals table → closure env → primitives → HO wrappers" + "Look up a global: closure env → globals table → primitives → HO forms" (let - ((globals (vm-globals-ref vm))) + ((closure-env (get (frame-closure frame) "vm-closure-env"))) (if - (has-key? globals name) - (get globals name) + (nil? closure-env) (let - ((closure-env (-> frame frame-closure closure-env))) + ((globals (vm-globals-ref vm))) (if - (nil? closure-env) + (has-key? globals name) + (get globals name) (cek-try (fn () (get-primitive name)) - (fn (e) (vm-resolve-ho-form vm name))) + (fn (e) (vm-resolve-ho-form vm name))))) + (let + ((found (env-walk closure-env name))) + (if + (nil? found) (let - ((found (env-walk closure-env name))) + ((globals (vm-globals-ref vm))) (if - (nil? found) + (has-key? globals name) + (get globals name) (cek-try (fn () (get-primitive name)) - (fn (e) (vm-resolve-ho-form vm name))) - found)))))))) + (fn (e) (vm-resolve-ho-form vm name))))) + found)))))) (define vm-resolve-ho-form (fn