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

Fix JIT compilation cascade + MCP robustness

Browse Source 
Three interacting JIT bugs caused infinite loops and server hangs:

1. _jit_compiling cascade: the re-entrancy flag was local to each
   binary's hook. When vm_call triggered JIT compilation internally,
   compiler functions got JIT-compiled during compilation, creating
   infinite cascades. Fix: shared _jit_compiling flag in sx_vm.ml,
   set in jit_compile_lambda itself.

2. call_closure always created new VMs: every HO primitive callback
   (for-each, map, filter) allocated a fresh VM. With 43K+ calls
   during compilation, this was the direct cause of hangs. Fix:
   call_closure_reuse reuses the active VM by isolating frames and
   running re-entrantly. VmSuspended is handled by merging frames
   for proper IO resumption.

3. vm_call for compiled Lambdas: OP_CALL dispatching to a Lambda
   with cached bytecode created a new VM instead of pushing a frame
   on the current one. Fix: push_closure_frame directly.

Additional MCP server fixes:
- Hot-reload: auto-execv when binary on disk is newer (no restart needed)
- Robust JSON: to_int_safe/to_int_or handle null, string, int params
- sx_summarise depth now optional (default 2)
- Per-request error handling (malformed JSON doesn't crash server)
- sx_test uses pre-built binary (skips dune rebuild overhead)
- Timed module loading for startup diagnostics

sx_server.ml fixes:
- Uses shared _jit_compiling flag
- Marks lambdas as jit_failed_sentinel on compile failure (no retry spam)
- call_closure_reuse with VmSuspended frame merging for IO support

Compiled compiler bytecode bug: deeply nested cond/case/let forms
(e.g. tw-resolve-style) cause the compiled compiler to loop.
Workaround: _jit_compiling guard prevents compiled function execution
during compilation. Compilation uses CEK (slower but correct).
Test suite: 3127/3127 passed.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
giles
2026-04-07 21:59:31 +00:00
parent 75130876c7
commit 03278c640d
3 changed files with 193 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
hosts/ocaml/bin/mcp_tree.ml
View File

@@ -6,6 +6,41 @@
open Sx_types open Sx_types
(* ------------------------------------------------------------------ *)
(* Hot-reload: re-exec ourselves when the binary has been rebuilt *)
(* ------------------------------------------------------------------ *)
let exe_path =
try Unix.readlink "/proc/self/exe"
with _ -> Sys.executable_name
let binary_mtime =
ref (try (Unix.stat exe_path).Unix.st_mtime with _ -> 0.0)
let check_hot_reload () =
try
let cur = (Unix.stat exe_path).Unix.st_mtime in
if cur > !binary_mtime then begin
Printf.eprintf "[mcp] Binary updated (%.0f -> %.0f), hot-reloading...\n%!" !binary_mtime cur;
Unix.execv exe_path [| exe_path |]
end
with _ -> ()
(* ------------------------------------------------------------------ *)
(* Robust JSON helpers — MCP clients send ints as strings or null *)
(* ------------------------------------------------------------------ *)
let to_int_safe json =
match json with
| `Int n -> Some n
| `Float f -> Some (int_of_float f)
| `String s -> (try Some (int_of_string s) with _ -> None)
| `Null -> None
| _ -> None
let to_int_or ~default json =
match to_int_safe json with Some n -> n | None -> default
(* ------------------------------------------------------------------ *) (* ------------------------------------------------------------------ *)
(* SX evaluator setup — minimal env for parser + tree-tools *) (* SX evaluator setup — minimal env for parser + tree-tools *)
(* ------------------------------------------------------------------ *) (* ------------------------------------------------------------------ *)
@@ -29,44 +64,44 @@ let load_sx_file e path =
(* JIT infrastructure — shared VM globals table, kept in sync via env_bind hook *) (* 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 _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 _jit_warned : (string, bool) Hashtbl.t = Hashtbl.create 32
let register_mcp_jit_hook () = let register_mcp_jit_hook () =
Sx_runtime._jit_try_call_fn := Some (fun f args -> Sx_runtime._jit_try_call_fn := Some (fun f args ->
match f with match f with
| Lambda l -> | Lambda l ->
let fn_name = match l.l_name with Some n -> n | None -> "" in
if fn_name <> "" && Hashtbl.mem _jit_warned fn_name then None
else
(match l.l_compiled with (match l.l_compiled with
| Some cl when not (Sx_vm.is_jit_failed cl) -> | Some cl when not (Sx_vm.is_jit_failed cl) ->
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref) (* Already compiled — run via VM. Skip during compilation. *)
if !(Sx_vm._jit_compiling) then None
else
(try Some (Sx_vm.call_closure_reuse cl args)
with e -> 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 if not (Hashtbl.mem _jit_warned fn_name) then begin
Hashtbl.replace _jit_warned fn_name true; Hashtbl.replace _jit_warned fn_name true;
Printf.eprintf "[mcp-jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e) Printf.eprintf "[mcp-jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e)
end; end;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None) None)
| Some _ -> None | Some _ -> None
| None -> | None ->
if !_jit_compiling then None (* Only block NEW compilations during _jit_compiling *)
else begin if !(Sx_vm._jit_compiling) then None
let fn_name = match l.l_name with Some n -> n | None -> "?" in else
if Hashtbl.mem _jit_warned fn_name then None let compiled = Sx_vm.jit_compile_lambda l _mcp_vm_globals in
else begin (match compiled with
_jit_compiling := true; | Some cl ->
let compiled = Sx_vm.jit_compile_lambda l _mcp_vm_globals in l.l_compiled <- Some cl;
_jit_compiling := false; (try Some (Sx_vm.call_closure_reuse cl args)
match compiled with with e ->
| Some cl -> Printf.eprintf "[mcp-jit] %s first-call fallback: %s\n%!" fn_name (Printexc.to_string e);
l.l_compiled <- Some cl; Hashtbl.replace _jit_warned fn_name true;
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref) l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
with e -> None)
Printf.eprintf "[mcp-jit] %s first-call fallback: %s\n%!" fn_name (Printexc.to_string e); | None -> None))
Hashtbl.replace _jit_warned fn_name true;
None)
| None -> None
end
end)
| _ -> None) | _ -> None)
let setup_env () = let setup_env () =
@@ -77,6 +112,13 @@ let setup_env () =
match v with match v with
| Thunk (body, closure_env) -> Sx_ref.eval_expr body (Env closure_env) | Thunk (body, closure_env) -> Sx_ref.eval_expr body (Env closure_env)
| other -> other); | other -> other);
(* Seed VM globals with all primitives as NativeFn values.
Without this, OP_CALL_PRIM in JIT-compiled bytecode can't find
primitives and falls back to CEK — which causes cascading failures
for stdlib functions like nth/first that override native versions. *)
Hashtbl.iter (fun name fn ->
Hashtbl.replace _mcp_vm_globals name (NativeFn (name, fn))
) Sx_primitives.primitives;
(* JIT: mirror root-env bindings into shared VM globals table *) (* JIT: mirror root-env bindings into shared VM globals table *)
Sx_types._env_bind_hook := Some (fun env name v -> Sx_types._env_bind_hook := Some (fun env name v ->
if env.parent = None then if env.parent = None then
@@ -304,51 +346,49 @@ let setup_env () =
| _ -> Nil); | _ -> Nil);
bind "deftype" (fun _args -> Nil); bind "deftype" (fun _args -> Nil);
bind "defeffect" (fun _args -> Nil); bind "defeffect" (fun _args -> Nil);
(* Load parser.sx for the SX-level sx-parse/sx-serialize *) (* Load SX modules with timing *)
let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
let lib_dir = try Sys.getenv "SX_LIB_DIR" with Not_found -> "lib" in let lib_dir = try Sys.getenv "SX_LIB_DIR" with Not_found -> "lib" in
(try load_sx_file e (Filename.concat spec_dir "parser.sx")
with exn -> Printf.eprintf "[mcp] Warning: parser.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load tree-tools *)
(try load_sx_file e (Filename.concat lib_dir "tree-tools.sx")
with exn -> Printf.eprintf "[mcp] Error: tree-tools.sx load failed: %s\n%!" (Printexc.to_string exn); exit 1);
(* Load signals — reactive signal primitives *)
(try load_sx_file e (Filename.concat spec_dir "signals.sx")
with exn -> Printf.eprintf "[mcp] Warning: signals.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Load render + adapter-html for render-to-html *)
let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in
(try load_sx_file e (Filename.concat spec_dir "render.sx") let t0 = Unix.gettimeofday () in
with exn -> Printf.eprintf "[mcp] Warning: render.sx load failed: %s\n%!" (Printexc.to_string exn)); let timed name f =
(try load_sx_file e (Filename.concat web_dir "adapter-html.sx") let t = Unix.gettimeofday () in
with exn -> Printf.eprintf "[mcp] Warning: adapter-html.sx load failed: %s\n%!" (Printexc.to_string exn)); f ();
(* Load harness *) Printf.eprintf "[mcp] %s: %.0fms\n%!" name ((Unix.gettimeofday () -. t) *. 1000.0)
(try load_sx_file e (Filename.concat spec_dir "harness.sx") in
with exn -> Printf.eprintf "[mcp] Warning: harness.sx load failed: %s\n%!" (Printexc.to_string exn)); timed "parser.sx" (fun () ->
(* Load eval-rules *) try load_sx_file e (Filename.concat spec_dir "parser.sx")
(try load_sx_file e (Filename.concat spec_dir "eval-rules.sx") with exn -> Printf.eprintf "[mcp] Warning: parser.sx load failed: %s\n%!" (Printexc.to_string exn));
with exn -> Printf.eprintf "[mcp] Warning: eval-rules.sx load failed: %s\n%!" (Printexc.to_string exn)); timed "tree-tools.sx" (fun () ->
(* Load signals — reactive signal primitives *) try load_sx_file e (Filename.concat lib_dir "tree-tools.sx")
(try load_sx_file e (Filename.concat spec_dir "signals.sx") with exn -> Printf.eprintf "[mcp] Error: tree-tools.sx load failed: %s\n%!" (Printexc.to_string exn); exit 1);
with exn -> Printf.eprintf "[mcp] Warning: signals.sx load failed: %s\n%!" (Printexc.to_string exn)); timed "signals.sx" (fun () ->
(* Load render + adapter-html for render-to-html *) try load_sx_file e (Filename.concat spec_dir "signals.sx")
let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found -> "web" in with exn -> Printf.eprintf "[mcp] Warning: signals.sx load failed: %s\n%!" (Printexc.to_string exn));
(try load_sx_file e (Filename.concat spec_dir "render.sx") timed "render.sx" (fun () ->
with exn -> Printf.eprintf "[mcp] Warning: render.sx load failed: %s\n%!" (Printexc.to_string exn)); try load_sx_file e (Filename.concat spec_dir "render.sx")
(try load_sx_file e (Filename.concat web_dir "adapter-html.sx") with exn -> Printf.eprintf "[mcp] Warning: render.sx load failed: %s\n%!" (Printexc.to_string exn));
with exn -> Printf.eprintf "[mcp] Warning: adapter-html.sx load failed: %s\n%!" (Printexc.to_string exn)); timed "adapter-html.sx" (fun () ->
(* Load render pipeline — native OCaml renderer + HTML tag bindings *) try load_sx_file e (Filename.concat web_dir "adapter-html.sx")
with exn -> Printf.eprintf "[mcp] Warning: adapter-html.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "harness.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "harness.sx")
with exn -> Printf.eprintf "[mcp] Warning: harness.sx load failed: %s\n%!" (Printexc.to_string exn));
timed "eval-rules.sx" (fun () ->
try load_sx_file e (Filename.concat spec_dir "eval-rules.sx")
with exn -> Printf.eprintf "[mcp] Warning: eval-rules.sx load failed: %s\n%!" (Printexc.to_string exn));
(* Render pipeline — native OCaml renderer + HTML tag bindings *)
Sx_render.setup_render_env e; Sx_render.setup_render_env e;
List.iter (fun tag -> List.iter (fun tag ->
ignore (Sx_types.env_bind e tag ignore (Sx_types.env_bind e tag
(NativeFn ("html:" ^ tag, fun args -> List (Symbol tag :: args)))) (NativeFn ("html:" ^ tag, fun args -> List (Symbol tag :: args))))
) Sx_render.html_tags; ) Sx_render.html_tags;
ignore (Sx_types.env_bind e "island?" (NativeFn ("island?", fun args -> match args with [Island _] -> Bool true | _ -> Bool false))); ignore (Sx_types.env_bind e "island?" (NativeFn ("island?", fun args -> match args with [Island _] -> Bool true | _ -> Bool false)));
(* Load compiler + enable JIT *) timed "compiler.sx+JIT" (fun () ->
(try load_sx_file e (Filename.concat lib_dir "compiler.sx"); try load_sx_file e (Filename.concat lib_dir "compiler.sx");
register_mcp_jit_hook (); 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));
with exn -> Printf.eprintf "[mcp] Warning: compiler.sx load failed (JIT disabled): %s\n%!" (Printexc.to_string exn)); Printf.eprintf "[mcp] Ready in %.0fms\n%!" ((Unix.gettimeofday () -. t0) *. 1000.0);
Printf.eprintf "[mcp] SX tree-tools + harness + eval-rules + render loaded\n%!";
env := e env := e
(* ------------------------------------------------------------------ *) (* ------------------------------------------------------------------ *)
@@ -616,9 +656,9 @@ let handle_tool name args =
let file = args |> member "file" |> to_string in let file = args |> member "file" |> to_string in
let tree, cst = parse_file_cst file in let tree, cst = parse_file_cst file in
let focus = args |> member "focus" |> to_string_option in let focus = args |> member "focus" |> to_string_option in
let max_depth = args |> member "max_depth" |> to_int_option in let max_depth = to_int_safe (args |> member "max_depth") in
let max_lines = args |> member "max_lines" |> to_int_option in let max_lines = to_int_safe (args |> member "max_lines") in
let offset = args |> member "offset" |> to_int_option |> Option.value ~default:0 in let offset = to_int_or ~default:0 (args |> member "offset") in
(match focus with (match focus with
| Some pattern -> | Some pattern ->
text_result (inject_cst_comments (value_to_string (call_sx "annotate-focused" [tree; String pattern])) (extract_cst_comments cst)) text_result (inject_cst_comments (value_to_string (call_sx "annotate-focused" [tree; String pattern])) (extract_cst_comments cst))
@@ -642,7 +682,7 @@ let handle_tool name args =
| "sx_summarise" -> | "sx_summarise" ->
let file = args |> member "file" |> to_string in let file = args |> member "file" |> to_string in
let tree, cst = parse_file_cst file in let tree, cst = parse_file_cst file in
let depth = args |> member "depth" |> to_int in let depth = to_int_or ~default:2 (args |> member "depth") in
text_result (inject_cst_comments (value_to_string (call_sx "summarise" [tree; Number (float_of_int depth)])) (extract_cst_comments cst)) text_result (inject_cst_comments (value_to_string (call_sx "summarise" [tree; Number (float_of_int depth)])) (extract_cst_comments cst))
| "sx_read_subtree" -> | "sx_read_subtree" ->
@@ -697,7 +737,7 @@ let handle_tool name args =
let file = args |> member "file" |> to_string in let file = args |> member "file" |> to_string in
let tree, cst = parse_file_cst file in let tree, cst = parse_file_cst file in
let path = resolve_path tree (args |> member "path" |> to_string) in let path = resolve_path tree (args |> member "path" |> to_string) in
let index = args |> member "index" |> to_int in let index = to_int_or ~default:0 (args |> member "index") in
let src = args |> member "new_source" |> to_string in let src = args |> member "new_source" |> to_string in
write_edit_cst file cst (call_sx "insert-child" [tree; path; Number (float_of_int index); String src]) write_edit_cst file cst (call_sx "insert-child" [tree; path; Number (float_of_int index); String src])
@@ -873,11 +913,18 @@ let handle_tool name args =
let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in let spec_dir = try Sys.getenv "SX_SPEC_DIR" with Not_found -> "spec" in
Filename.dirname spec_dir Filename.dirname spec_dir
in in
let timeout = args |> member "timeout" |> to_int_option |> Option.value ~default:300 in let timeout = to_int_or ~default:300 (args |> member "timeout") in
let cmd = match host with let cmd = match host with
| "ocaml" -> | "ocaml" ->
Printf.sprintf "cd %s/hosts/ocaml && eval $(opam env 2>/dev/null) && timeout %d dune exec bin/run_tests.exe%s 2>&1" (* Use pre-built binary directly — avoids dune rebuild delay.
project_dir timeout (if full then " -- --full" else "") Falls back to dune exec if the binary doesn't exist. *)
let exe = Printf.sprintf "%s/hosts/ocaml/_build/default/bin/run_tests.exe" project_dir in
if Sys.file_exists exe then
Printf.sprintf "cd %s/hosts/ocaml && timeout %d %s%s 2>&1"
project_dir timeout exe (if full then " --full" else "")
else
Printf.sprintf "cd %s/hosts/ocaml && eval $(opam env 2>/dev/null) && timeout %d dune exec bin/run_tests.exe%s 2>&1"
project_dir timeout (if full then " -- --full" else "")
| "js" | _ -> | "js" | _ ->
Printf.sprintf "cd %s && timeout %d node hosts/javascript/run_tests.js%s 2>&1" Printf.sprintf "cd %s && timeout %d node hosts/javascript/run_tests.js%s 2>&1"
project_dir timeout (if full then " --full" else "") project_dir timeout (if full then " --full" else "")
@@ -1827,7 +1874,7 @@ let handle_tool name args =
| "sx_trace" -> | "sx_trace" ->
let expr_str = args |> member "expr" |> to_string in let expr_str = args |> member "expr" |> to_string in
let max_steps = (try args |> member "max_steps" |> to_int with _ -> 200) in let max_steps = to_int_or ~default:200 (args |> member "max_steps") in
let file = try Some (args |> member "file" |> to_string) with _ -> None in let file = try Some (args |> member "file" |> to_string) with _ -> None in
let components_only = (try args |> member "components_only" |> to_bool with _ -> false) in let components_only = (try args |> member "components_only" |> to_bool with _ -> false) in
let e = !env in let e = !env in
@@ -2364,7 +2411,7 @@ let tool_definitions = `List [
("offset", `Assoc [("type", `String "integer"); ("description", `String "Line offset for pagination (default 0)")])] ("offset", `Assoc [("type", `String "integer"); ("description", `String "Line offset for pagination (default 0)")])]
["file"]; ["file"];
tool "sx_summarise" "Folded structural overview of an .sx file. Use to orient before drilling into a region." tool "sx_summarise" "Folded structural overview of an .sx file. Use to orient before drilling into a region."
[file_prop; ("depth", `Assoc [("type", `String "integer"); ("description", `String "Max depth (0=heads only)")])] ["file"; "depth"]; [file_prop; ("depth", `Assoc [("type", `String "integer"); ("description", `String "Max depth (0=heads only, default 2)")])] ["file"];
tool "sx_read_subtree" "Expand a specific subtree by path. Use after summarise to drill in." tool "sx_read_subtree" "Expand a specific subtree by path. Use after summarise to drill in."
[file_prop; path_prop] ["file"; "path"]; [file_prop; path_prop] ["file"; "path"];
tool "sx_get_context" "Show enclosing chain from root to a target node." tool "sx_get_context" "Show enclosing chain from root to a target node."
@@ -2542,21 +2589,28 @@ let () =
try while true do try while true do
let line = input_line stdin in let line = input_line stdin in
if String.length line > 0 then begin if String.length line > 0 then begin
let json = Yojson.Safe.from_string line in try
let open Yojson.Safe.Util in let json = Yojson.Safe.from_string line in
let meth = json |> member "method" |> to_string_option |> Option.value ~default:"" in let open Yojson.Safe.Util in
let params = json |> member "params" in let meth = json |> member "method" |> to_string_option |> Option.value ~default:"" in
let id = json |> member "id" in let params = json |> member "params" in
let result = dispatch meth params in let id = json |> member "id" in
if id <> `Null then begin let result =
let resp = `Assoc [ try dispatch meth params
("jsonrpc", `String "2.0"); with e -> error_result ("Error: " ^ Printexc.to_string e)
("id", id); in
("result", result)] in if id <> `Null then begin
print_string (Yojson.Safe.to_string resp); let resp = `Assoc [
print_char '\n'; ("jsonrpc", `String "2.0");
flush stdout ("id", id);
end ("result", result)] in
print_string (Yojson.Safe.to_string resp);
print_char '\n';
flush stdout;
check_hot_reload ()
end
with e ->
Printf.eprintf "[mcp] Request error: %s\n%!" (Printexc.to_string e)
end end
done done
with End_of_file -> () with End_of_file -> ()

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

@@ -908,7 +908,7 @@ let sx_render_to_html expr env =
into the kernel env. The hook handles both cached execution (bytecode into the kernel env. The hook handles both cached execution (bytecode
already compiled) and first-call compilation (invoke compiler.sx via already compiled) and first-call compilation (invoke compiler.sx via
CEK, cache result). cek_call checks this before CEK dispatch. *) CEK, cache result). cek_call checks this before CEK dispatch. *)
let _jit_compiling = ref false (* re-entrancy guard *) (* Re-entrancy guard lives in Sx_vm._jit_compiling (shared with vm_call path) *)
(* JIT compilation is lazy-only: every named lambda gets one compile (* JIT compilation is lazy-only: every named lambda gets one compile
attempt on first call. Failures are sentineled (never retried). *) attempt on first call. Failures are sentineled (never retried). *)
@@ -936,9 +936,11 @@ let register_jit_hook env =
| Lambda l -> | Lambda l ->
(match l.l_compiled with (match l.l_compiled with
| Some cl when not (Sx_vm.is_jit_failed cl) -> | Some cl when not (Sx_vm.is_jit_failed cl) ->
(* Cached bytecode — run on VM, fall back to CEK on runtime error. (* Cached bytecode — run via VM. Skip during compilation to avoid
Log once per function name, then stay quiet. Don't disable. *) compiled-compiler bytecode bugs on complex nested forms. *)
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref) if !(Sx_vm._jit_compiling) then None
else
(try Some (Sx_vm.call_closure_reuse cl args)
with with
| Sx_vm.VmSuspended (request, saved_vm) -> | Sx_vm.VmSuspended (request, saved_vm) ->
Some (make_vm_suspend_marker request saved_vm) Some (make_vm_suspend_marker request saved_vm)
@@ -948,33 +950,34 @@ let register_jit_hook env =
Hashtbl.replace _jit_warned fn_name true; Hashtbl.replace _jit_warned fn_name true;
Printf.eprintf "[jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e) Printf.eprintf "[jit] %s runtime fallback to CEK: %s\n%!" fn_name (Printexc.to_string e)
end; end;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None) None)
| Some _ -> None (* compile failed or disabled — CEK handles *) | Some _ -> None
| None -> | None ->
(* Only block NEW compilations during _jit_compiling, not execution *)
let fn_name = match l.l_name with Some n -> n | None -> "?" in let fn_name = match l.l_name with Some n -> n | None -> "?" in
if !_jit_compiling then None if !(Sx_vm._jit_compiling) then None
else if Hashtbl.mem _jit_warned fn_name then None else if Hashtbl.mem _jit_warned fn_name then None
else begin else begin
_jit_compiling := true;
let t0 = Unix.gettimeofday () in let t0 = Unix.gettimeofday () in
let compiled = Sx_vm.jit_compile_lambda l (env_to_vm_globals env) in let compiled = Sx_vm.jit_compile_lambda l (env_to_vm_globals env) in
let dt = Unix.gettimeofday () -. t0 in let dt = Unix.gettimeofday () -. t0 in
_jit_compiling := false; Printf.eprintf "[jit] %s compile in %.3fs\n%!" fn_name dt;
if dt > 0.5 || (match compiled with None -> true | _ -> false) then
Printf.eprintf "[jit] %s compile %s in %.3fs\n%!"
fn_name (match compiled with Some _ -> "OK" | None -> "FAIL") dt;
match compiled with match compiled with
| Some cl -> | Some cl ->
l.l_compiled <- Some cl; l.l_compiled <- Some cl;
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref) (try Some (Sx_vm.call_closure_reuse cl args)
with with
| Sx_vm.VmSuspended (request, saved_vm) -> | Sx_vm.VmSuspended (request, saved_vm) ->
Some (make_vm_suspend_marker request saved_vm) Some (make_vm_suspend_marker request saved_vm)
| e -> | e ->
Printf.eprintf "[jit] %s first-call fallback to CEK: %s\n%!" fn_name (Printexc.to_string e); Printf.eprintf "[jit] %s first-call fallback to CEK: %s\n%!" fn_name (Printexc.to_string e);
Hashtbl.replace _jit_warned fn_name true; Hashtbl.replace _jit_warned fn_name true;
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None) None)
| None -> None | None ->
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
None
end) end)
| _ -> None) | _ -> None)

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

@@ -133,6 +133,11 @@ let vm_report_counters () =
Printf.eprintf "[vm-perf] insns=%d calls=%d cek_fallbacks=%d comp_jit=%d comp_cek=%d\n%!" Printf.eprintf "[vm-perf] insns=%d calls=%d cek_fallbacks=%d comp_jit=%d comp_cek=%d\n%!"
!_vm_insn_count !_vm_call_count !_vm_cek_count !_vm_comp_jit_count !_vm_comp_cek_count !_vm_insn_count !_vm_call_count !_vm_cek_count !_vm_comp_jit_count !_vm_comp_cek_count
(** Global flag: true while a JIT compilation is in progress.
Prevents the JIT hook from intercepting calls during compilation,
which would cause infinite cascades (compiling the compiler). *)
let _jit_compiling = ref false
(** Push a VM closure frame onto the current VM — no new VM allocation. (** Push a VM closure frame onto the current VM — no new VM allocation.
This is the fast path for intra-VM closure calls. *) This is the fast path for intra-VM closure calls. *)
let push_closure_frame vm cl args = let push_closure_frame vm cl args =
@@ -228,9 +233,30 @@ let rec call_closure cl args globals =
(** Call a VmClosure on the active VM if one exists, otherwise create a new one. (** Call a VmClosure on the active VM if one exists, otherwise create a new one.
This is the path used by HO primitives (map, filter, for-each, some) so This is the path used by HO primitives (map, filter, for-each, some) so
callbacks can access upvalues that reference the calling VM's state. *) callbacks run on the same VM, avoiding per-call VM allocation overhead. *)
and call_closure_reuse cl args = and call_closure_reuse cl args =
call_closure cl args cl.vm_env_ref match !_active_vm with
| Some vm ->
let saved_sp = vm.sp in
push_closure_frame vm cl args;
let saved_frames = List.tl vm.frames in
vm.frames <- [List.hd vm.frames];
(try run vm
with
| VmSuspended _ as e ->
(* IO suspension: merge remaining callback frames with caller frames
so the VM can be properly resumed. When resumed, it finishes the
callback then returns to the caller's frames. *)
vm.frames <- vm.frames @ saved_frames;
raise e
| e ->
vm.frames <- saved_frames;
vm.sp <- saved_sp;
raise e);
vm.frames <- saved_frames;
pop vm
| None ->
call_closure cl args cl.vm_env_ref
(** Call a value as a function — dispatch by type. (** Call a value as a function — dispatch by type.
VmClosure: pushes frame on current VM (fast intra-VM path). VmClosure: pushes frame on current VM (fast intra-VM path).
@@ -247,25 +273,18 @@ and vm_call vm f args =
| Lambda l -> | Lambda l ->
(match l.l_compiled with (match l.l_compiled with
| Some cl when not (is_jit_failed cl) -> | Some cl when not (is_jit_failed cl) ->
(* Cached bytecode — run on VM using the closure's captured env, (* Cached bytecode — push frame on current VM *)
not the caller's globals. Closure vars were merged at compile time. *) push_closure_frame vm cl args
(try push vm (call_closure cl args cl.vm_env_ref)
with _e ->
(* Fallback to CEK — suspension-aware *)
push vm (cek_call_or_suspend vm f (List args)))
| Some _ -> | Some _ ->
(* Compile failed — CEK, suspension-aware *)
push vm (cek_call_or_suspend vm f (List args)) push vm (cek_call_or_suspend vm f (List args))
| None -> | None ->
if l.l_name <> None if l.l_name <> None
then begin then begin
(* Pre-mark before compile attempt to prevent re-entrancy *)
l.l_compiled <- Some jit_failed_sentinel; l.l_compiled <- Some jit_failed_sentinel;
match !jit_compile_ref l vm.globals with match !jit_compile_ref l vm.globals with
| Some cl -> | Some cl ->
l.l_compiled <- Some cl; l.l_compiled <- Some cl;
(try push vm (call_closure cl args cl.vm_env_ref) push_closure_frame vm cl args
with _e -> push vm (cek_call_or_suspend vm f (List args)))
| None -> | None ->
push vm (cek_call_or_suspend vm f (List args)) push vm (cek_call_or_suspend vm f (List args))
end end
@@ -784,9 +803,14 @@ let execute_module_safe code globals =
record so subsequent calls go straight to the VM. *) record so subsequent calls go straight to the VM. *)
let jit_compile_lambda (l : lambda) globals = let jit_compile_lambda (l : lambda) globals =
let fn_name = match l.l_name with Some n -> n | None -> "<anon>" in let fn_name = match l.l_name with Some n -> n | None -> "<anon>" in
if !_jit_compiling then (
(* Already compiling — prevent cascade. The CEK will handle this call. *)
None
) else
try try
_jit_compiling := true;
let compile_fn = try Hashtbl.find globals "compile" let compile_fn = try Hashtbl.find globals "compile"
with Not_found -> raise (Eval_error "JIT: compiler not loaded") in with Not_found -> (_jit_compiling := false; raise (Eval_error "JIT: compiler not loaded")) in
(* Reconstruct the (fn (params) body) form so the compiler produces (* Reconstruct the (fn (params) body) form so the compiler produces
a proper closure. l.l_body is the inner body; we need the full a proper closure. l.l_body is the inner body; we need the full
function form with params so the compiled code binds them. *) function form with params so the compiled code binds them. *)
@@ -800,12 +824,7 @@ let jit_compile_lambda (l : lambda) globals =
let compile_env = Sx_types.env_extend (Sx_types.make_env ()) in let compile_env = Sx_types.env_extend (Sx_types.make_env ()) in
Hashtbl.iter (fun k v -> Hashtbl.replace compile_env.bindings (Sx_types.intern k) v) globals; Hashtbl.iter (fun k v -> Hashtbl.replace compile_env.bindings (Sx_types.intern k) v) globals;
let result = Sx_ref.eval_expr (List [Symbol "compile"; quoted]) (Env compile_env) in let result = Sx_ref.eval_expr (List [Symbol "compile"; quoted]) (Env compile_env) in
(* Closure vars are accessible via vm_closure_env (set on the VmClosure _jit_compiling := false;
at line ~617). OP_GLOBAL_GET falls back to vm_closure_env when vars
aren't in globals. No injection into the shared globals table —
that would break closure isolation for factory functions like
make-page-fn where multiple closures capture different values
for the same variable names. *)
let effective_globals = globals in let effective_globals = globals in
(match result with (match result with
| Dict d when Hashtbl.mem d "bytecode" -> | Dict d when Hashtbl.mem d "bytecode" ->
@@ -821,21 +840,13 @@ let jit_compile_lambda (l : lambda) globals =
else begin else begin
Printf.eprintf "[jit] FAIL %s: closure index %d out of bounds (pool=%d)\n%!" Printf.eprintf "[jit] FAIL %s: closure index %d out of bounds (pool=%d)\n%!"
fn_name idx (Array.length outer_code.vc_constants); fn_name idx (Array.length outer_code.vc_constants);
None None
end end
end else begin end else begin
(* Not a closure — constant expression, alias, or simple computation.
Execute the bytecode as a module to get the value, then wrap
as a NativeFn if it's callable (so the CEK can dispatch to it). *)
(try (try
let value = execute_module outer_code globals in let value = execute_module outer_code globals in
Printf.eprintf "[jit] RESOLVED %s: %s (bc[0]=%d)\n%!" 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); fn_name (type_of value) (if Array.length bc > 0 then bc.(0) else -1);
(* 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 None
with _ -> with _ ->
Printf.eprintf "[jit] SKIP %s: non-closure execution failed (bc[0]=%d, len=%d)\n%!" Printf.eprintf "[jit] SKIP %s: non-closure execution failed (bc[0]=%d, len=%d)\n%!"
@@ -846,12 +857,13 @@ let jit_compile_lambda (l : lambda) globals =
Printf.eprintf "[jit] FAIL %s: compiler returned %s\n%!" fn_name (type_of result); Printf.eprintf "[jit] FAIL %s: compiler returned %s\n%!" fn_name (type_of result);
None) None)
with e -> with e ->
_jit_compiling := false;
Printf.eprintf "[jit] FAIL %s: %s\n%!" fn_name (Printexc.to_string e); Printf.eprintf "[jit] FAIL %s: %s\n%!" fn_name (Printexc.to_string e);
None None
(* Wire up forward references *) (* Wire up forward references *)
let () = jit_compile_ref := jit_compile_lambda let () = jit_compile_ref := jit_compile_lambda
let () = _vm_call_closure_ref := (fun cl args -> call_closure cl args cl.vm_env_ref) let () = _vm_call_closure_ref := (fun cl args -> call_closure_reuse cl args)
(** {1 Debugging / introspection} *) (** {1 Debugging / introspection} *)