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

Lazy JIT compilation: lambdas compile to bytecode on first call

Browse Source 
Replace AOT adapter compilation with lazy JIT — each named lambda is
compiled to VM bytecode on first call, cached in l_compiled field for
subsequent calls. Compilation failures fall back to CEK gracefully.

VM types (vm_code, vm_upvalue_cell, vm_closure) moved to sx_types.ml
mutual recursion block. Lambda and Component records gain mutable
l_compiled/c_compiled cache fields. jit_compile_lambda in sx_vm.ml
wraps body as (fn (params) body), invokes spec/compiler.sx via CEK,
extracts inner closure from OP_CLOSURE constant.

JIT hooks in both paths:
- vm_call: Lambda calls from compiled VM code
- continue_with_call: Lambda calls from CEK step loop (injected by
  bootstrap.py post-processing)

Pre-mark sentinel prevents re-entrancy (compile function itself was
hanging when JIT'd mid-compilation). VM execution errors caught and
fall back to CEK with sentinel marking.

Also: add kbd/samp/var to HTML_TAGS, rebuild sx-browser.js, add page
URL to sx-page-full-py timing log.

Performance: first page 28s (JIT compiles 17 functions), subsequent
pages 0.31s home / 0.71s wittgenstein (was 2.3s). All 1945 tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
giles
2026-03-23 08:18:44 +00:00
parent 7628659854
commit 318c818728
10 changed files with 296 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -698,96 +698,60 @@ let make_server_env () =
(* ====================================================================== *)
(* VM adapter — compiled aser functions in isolated globals *)
(* JIT hook registration *)
(* ====================================================================== *)
(** Compiled adapter globals — separate from kernel env.
Contains compiled aser functions + reads from kernel env for
components, helpers, and other runtime bindings. *)
let vm_adapter_globals : (string, value) Hashtbl.t option ref = ref None
(** Register the JIT call hook. Called once after the compiler is loaded
into the kernel env. The hook handles both cached execution (bytecode
already compiled) and first-call compilation (invoke compiler.sx via
CEK, cache result). cek_call checks this before CEK dispatch. *)
let _jit_compiling = ref false (* re-entrancy guard *)
(** Compile adapter-sx.sx and store in vm_adapter_globals.
Called from vm-compile-adapter command. *)
let compile_adapter env =
if not (Hashtbl.mem env.bindings "compile") then
raise (Eval_error "compiler not loaded")
else begin
let compile_fn = Hashtbl.find env.bindings "compile" in
(* Find and parse adapter-sx.sx *)
let web_dir = try Sys.getenv "SX_WEB_DIR" with Not_found ->
try Filename.concat (Sys.getenv "SX_SPEC_DIR") "../web"
with Not_found -> "web" in
let adapter_path = Filename.concat web_dir "adapter-sx.sx" in
if not (Sys.file_exists adapter_path) then
raise (Eval_error ("adapter-sx.sx not found: " ^ adapter_path));
let exprs = Sx_parser.parse_file adapter_path in
(* Compile each define's body *)
let globals = Hashtbl.create 64 in
(* Seed with kernel env for component/helper lookups *)
Hashtbl.iter (fun k v -> Hashtbl.replace globals k v) env.bindings;
let compiled = ref 0 in
List.iter (fun expr ->
match expr with
| List (Symbol "define" :: Symbol name :: rest) ->
(* Find the body — skip :effects annotations *)
let rec find_body = function
| Keyword _ :: _ :: rest -> find_body rest
| body :: _ -> body
| [] -> Nil
in
let body = find_body rest in
(try
let quoted = List [Symbol "quote"; body] in
let result = Sx_ref.eval_expr (List [compile_fn; quoted]) (Env env) in
match result with
| Dict d when Hashtbl.mem d "bytecode" ->
let outer_code = Sx_vm.code_from_value result in
Printf.eprintf "[vm] %s: outer bc=%d consts=%d inner_type=%s\n%!"
name (Array.length outer_code.Sx_vm.bytecode)
(Array.length outer_code.Sx_vm.constants)
(if Array.length outer_code.Sx_vm.constants > 0 then
type_of outer_code.Sx_vm.constants.(0) else "empty");
let bc = outer_code.Sx_vm.bytecode in
if Array.length bc >= 4 && bc.(0) = 51 then begin
(* The compiled define body is (fn ...) which compiles to
OP_CLOSURE + [upvalue descriptors] + OP_RETURN.
Extract the inner code object from constants[idx]. *)
let idx = bc.(1) lor (bc.(2) lsl 8) in
let code =
if idx < Array.length outer_code.Sx_vm.constants then begin
let inner_val = outer_code.Sx_vm.constants.(idx) in
try Sx_vm.code_from_value inner_val
with e ->
Printf.eprintf "[vm] inner code_from_value failed for %s: %s\n%!"
name (Printexc.to_string e);
raise e
end else outer_code
in
let cl = { Sx_vm.code; upvalues = [||]; name = Some name;
env_ref = globals } in
Hashtbl.replace globals name
(NativeFn ("vm:" ^ name, fun args ->
Sx_vm.call_closure cl args globals));
incr compiled
end else begin
(* Not a lambda — constant expression (e.g. (list ...)).
Execute once and store the resulting value directly. *)
let value = Sx_vm.execute_module outer_code globals in
Hashtbl.replace globals name value;
Printf.eprintf "[vm] %s: constant (type=%s)\n%!" name (type_of value);
incr compiled
end
| _ -> () (* non-dict result — skip *)
with e ->
Printf.eprintf "[vm] FAIL adapter %s: %s\n%!" name (Printexc.to_string e))
| _ ->
(* Non-define expression — evaluate on CEK to set up constants *)
(try ignore (Sx_ref.eval_expr expr (Env env)) with _ -> ())
) exprs;
vm_adapter_globals := Some globals;
Printf.eprintf "[vm] Compiled adapter: %d functions\n%!" !compiled
end
let register_jit_hook env =
Sx_ref.jit_call_hook := Some (fun f args ->
match f with
| Lambda l ->
(match l.l_compiled with
| Some cl when not (Sx_vm.is_jit_failed cl) ->
(* Cached bytecode — execute on VM, fall back to CEK on error *)
(try Some (Sx_vm.call_closure cl args cl.vm_env_ref)
with _ -> l.l_compiled <- Some Sx_vm.jit_failed_sentinel; None)
| Some _ -> None (* failed sentinel *)
| None ->
(* Don't try to compile while already compiling (prevents
infinite recursion: compile calls lambdas internally) *)
if !_jit_compiling then None
else begin
let fn_name = match l.l_name with Some n -> n | None -> "?" in
begin
(* Mark as tried BEFORE compiling — prevents other calls to
the same lambda from starting redundant compilations while
this one is running. If compilation succeeds, overwrite. *)
l.l_compiled <- Some Sx_vm.jit_failed_sentinel;
Printf.eprintf "[jit-hook] compiling %s (body size ~%d)...\n%!"
fn_name (String.length (inspect l.l_body));
_jit_compiling := true;
let t0 = Unix.gettimeofday () in
let compiled = Sx_vm.jit_compile_lambda l env.bindings in
let dt = Unix.gettimeofday () -. t0 in
_jit_compiling := false;
Printf.eprintf "[jit-hook] %s compile %s in %.3fs\n%!"
fn_name (match compiled with Some _ -> "OK" | None -> "FAIL") dt;
match compiled with
| Some cl ->
l.l_compiled <- Some cl;
Printf.eprintf "[jit-hook] executing %s on VM...\n%!" fn_name;
(try
let r = Sx_vm.call_closure cl args cl.vm_env_ref in
Printf.eprintf "[jit-hook] %s execution OK\n%!" fn_name;
Some r
with e ->
Printf.eprintf "[jit-hook] %s VM FAIL: %s\n%!" fn_name (Printexc.to_string e);
l.l_compiled <- Some Sx_vm.jit_failed_sentinel; None)
| None -> None
end end)
| _ -> None)
(* ====================================================================== *)
@@ -901,17 +865,15 @@ let rec dispatch env cmd =
| exn -> send_error (Printexc.to_string exn))
| List [Symbol "vm-compile-adapter"] ->
(* Compile adapter-sx.sx to VM bytecode with isolated globals *)
(try
compile_adapter env;
send_ok ()
with
| Eval_error msg -> send_error msg
| exn -> send_error (Printexc.to_string exn))
(* Legacy command — JIT hook is now registered at startup.
Kept for backward compatibility with ocaml_bridge.py. *)
register_jit_hook env;
Printf.eprintf "[jit] JIT hook registered (lazy compilation active)\n%!";
send_ok ()
| List [Symbol "aser-slot"; String src] ->
(* Expand ALL components server-side. Uses batch IO mode.
Routes through VM if adapter is compiled, else CEK. *)
Calls aser via CEK — the JIT hook compiles it on first call. *)
(try
let exprs = Sx_parser.parse_all src in
let expr = match exprs with
@@ -925,22 +887,14 @@ let rec dispatch env cmd =
let t0 = Unix.gettimeofday () in
let expand_fn = NativeFn ("expand-components?", fun _args -> Bool true) in
ignore (env_bind env "expand-components?" expand_fn);
let result = match !vm_adapter_globals with
| Some globals ->
Hashtbl.replace globals "expand-components?" expand_fn;
let aser_fn = try Hashtbl.find globals "aser"
with Not_found -> raise (Eval_error "VM: aser not compiled") in
let r = match aser_fn with
| NativeFn (_, fn) -> fn [expr; Env env]
| _ -> raise (Eval_error "VM: aser not a function")
in
Hashtbl.remove globals "expand-components?";
r
| None ->
let call = List [Symbol "aser";
List [Symbol "quote"; expr];
Env env] in
Sx_ref.eval_expr call (Env env)
Printf.eprintf "[aser-slot] starting aser eval...\n%!";
let result =
let call = List [Symbol "aser";
List [Symbol "quote"; expr];
Env env] in
let r = Sx_ref.eval_expr call (Env env) in
Printf.eprintf "[aser-slot] aser eval returned\n%!";
r
in
let t1 = Unix.gettimeofday () in
io_batch_mode := false;
@@ -990,22 +944,14 @@ let rec dispatch env cmd =
let t0 = Unix.gettimeofday () in
let expand_fn = NativeFn ("expand-components?", fun _args -> Bool true) in
ignore (env_bind env "expand-components?" expand_fn);
let body_result = match !vm_adapter_globals with
| Some globals ->
Hashtbl.replace globals "expand-components?" expand_fn;
let aser_fn = try Hashtbl.find globals "aser"
with Not_found -> raise (Eval_error "VM: aser not compiled") in
let r = match aser_fn with
| NativeFn (_, fn) -> fn [expr; Env env]
| _ -> raise (Eval_error "VM: aser not a function")
in
Hashtbl.remove globals "expand-components?";
r
| None ->
Printf.eprintf "[sx-page-full] starting aser eval...\n%!";
let body_result =
let call = List [Symbol "aser";
List [Symbol "quote"; expr];
Env env] in
Sx_ref.eval_expr call (Env env)
let r = Sx_ref.eval_expr call (Env env) in
Printf.eprintf "[sx-page-full] aser eval returned\n%!";
r
in
let t1 = Unix.gettimeofday () in
io_batch_mode := false;
@@ -1117,8 +1063,8 @@ let rec dispatch env cmd =
let fn = NativeFn ("vm:" ^ name, fun args ->
try
Sx_vm.call_closure
{ Sx_vm.code; upvalues = [||]; name = lam.l_name;
env_ref = live_env }
{ vm_code = code; vm_upvalues = [||]; vm_name = lam.l_name;
vm_env_ref = live_env }
args live_env
with
| _ ->