(** SX bytecode VM — stack-based interpreter.

    Executes bytecode produced by compiler.sx.
    Designed for speed: array-based stack, direct dispatch,
    no allocation per step (unlike the CEK machine).

    This is the platform-native execution engine. The same bytecode
    runs on all platforms (OCaml, JS, WASM).

    VM types (vm_code, vm_upvalue_cell, vm_closure) are defined in
    sx_types.ml to share the mutual recursion block with [value]. *)

open Sx_types

(** Call frame — one per function invocation. *)
type frame = {
  closure   : vm_closure;
  mutable ip : int;
  base      : int;  (* base index in value stack for locals *)
  local_cells : (int, vm_upvalue_cell) Hashtbl.t;  (* slot → shared cell for captured locals *)
}

(** VM state. *)
type vm = {
  mutable stack  : value array;
  mutable sp     : int;
  mutable frames : frame list;
  globals        : (string, value) Hashtbl.t;  (* live reference to kernel env *)
}

(** Forward reference for JIT compilation — set after definition. *)
let jit_compile_ref : (lambda -> (string, value) Hashtbl.t -> vm_closure option) ref =
  ref (fun _ _ -> None)

(** Sentinel closure indicating JIT compilation was attempted and failed.
    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_closure_env = None
}

let is_jit_failed cl = cl.vm_code.vc_arity = -1

let create globals =
  { stack = Array.make 4096 Nil; sp = 0; frames = []; globals }

(** Stack ops — inlined for speed. *)
let push vm v =
  if vm.sp >= Array.length vm.stack then begin
    let ns = Array.make (vm.sp * 2) Nil in
    Array.blit vm.stack 0 ns 0 vm.sp;
    vm.stack <- ns
  end;
  vm.stack.(vm.sp) <- v;
  vm.sp <- vm.sp + 1

let[@inline] pop vm =
  vm.sp <- vm.sp - 1;
  vm.stack.(vm.sp)

let[@inline] peek vm = vm.stack.(vm.sp - 1)

(** Read operands. *)
let[@inline] read_u8 f =
  let v = f.closure.vm_code.vc_bytecode.(f.ip) in
  f.ip <- f.ip + 1; v

let[@inline] read_u16 f =
  let lo = f.closure.vm_code.vc_bytecode.(f.ip) in
  let hi = f.closure.vm_code.vc_bytecode.(f.ip + 1) in
  f.ip <- f.ip + 2;
  lo lor (hi lsl 8)

let[@inline] read_i16 f =
  let v = read_u16 f in
  if v >= 32768 then v - 65536 else v

(** Wrap a VM closure as an SX value (NativeFn). *)
let closure_to_value cl =
  NativeFn ("vm:" ^ (match cl.vm_name with Some n -> n | None -> "anon"),
    fun args -> raise (Eval_error ("VM_CLOSURE_CALL:" ^ String.concat "," (List.map Sx_runtime.value_to_str args))))
  (* Placeholder — actual calls go through vm_call below *)

let _vm_insn_count = ref 0
let _vm_call_count = ref 0
let _vm_cek_count = ref 0
let vm_reset_counters () = _vm_insn_count := 0; _vm_call_count := 0; _vm_cek_count := 0
let vm_report_counters () =
  Printf.eprintf "[vm-perf] insns=%d calls=%d cek_fallbacks=%d\n%!"
    !_vm_insn_count !_vm_call_count !_vm_cek_count

(** Push a VM closure frame onto the current VM — no new VM allocation.
    This is the fast path for intra-VM closure calls. *)
let push_closure_frame vm cl args =
  let frame = { closure = cl; ip = 0; base = vm.sp; local_cells = Hashtbl.create 4 } in
  List.iter (fun a -> push vm a) args;
  for _ = List.length args to cl.vm_code.vc_locals - 1 do push vm Nil done;
  vm.frames <- frame :: vm.frames

(** Convert compiler output (SX dict) to a vm_code object. *)
let code_from_value v =
  match v with
  | Dict d ->
    let bc_list = match Hashtbl.find_opt d "bytecode" with
      | Some (List l | ListRef { contents = l }) ->
        Array.of_list (List.map (fun x -> match x with Number n -> int_of_float n | _ -> 0) l)
      | _ -> [||]
    in
    let entries = match Hashtbl.find_opt d "constants" with
      | Some (List l | ListRef { contents = l }) -> Array.of_list l
      | _ -> [||]
    in
    let constants = Array.map (fun entry ->
      match entry with
      | Dict ed when Hashtbl.mem ed "bytecode" -> entry  (* nested code — convert lazily *)
      | _ -> entry
    ) entries in
    let arity = match Hashtbl.find_opt d "arity" with
      | Some (Number n) -> int_of_float n | _ -> 0
    in
    { vc_arity = arity; vc_locals = arity + 16; vc_bytecode = bc_list; vc_constants = constants }
  | _ -> { vc_arity = 0; vc_locals = 16; vc_bytecode = [||]; vc_constants = [||] }

(** Execute a closure with arguments — creates a fresh VM.
    Used for entry points: JIT Lambda calls, module execution, cross-boundary. *)
let rec call_closure cl args globals =
  incr _vm_call_count;
  let vm = create globals in
  push_closure_frame vm cl args;
  (try run vm with e -> raise e);
  pop vm

(** Call a value as a function — dispatch by type.
    VmClosure: pushes frame on current VM (fast intra-VM path).
    Lambda: tries JIT then falls back to CEK.
    NativeFn: calls directly. *)
and vm_call vm f args =
  match f with
  | VmClosure cl ->
    (* Fast path: push frame on current VM — no allocation, enables TCO *)
    push_closure_frame vm cl args
  | NativeFn (_name, fn) ->
    let result = fn args in
    push vm result
  | Lambda l ->
    (match l.l_compiled with
     | Some cl when not (is_jit_failed cl) ->
       (* Cached bytecode — run on VM, fall back to CEK on runtime error *)
       (try push vm (call_closure cl args vm.globals)
        with _ -> push vm (Sx_ref.cek_call f (List args)))
     | Some _ ->
       (* Compile failed — CEK *)
       push vm (Sx_ref.cek_call f (List args))
     | None ->
       if l.l_name <> None
       then begin
         (* Pre-mark before compile attempt to prevent re-entrancy *)
         l.l_compiled <- Some jit_failed_sentinel;
         match !jit_compile_ref l vm.globals with
         | Some cl ->
           l.l_compiled <- Some cl;
           (try push vm (call_closure cl args vm.globals)
            with _ ->
              l.l_compiled <- Some jit_failed_sentinel;
              push vm (Sx_ref.cek_call f (List args)))
         | None ->
           push vm (Sx_ref.cek_call f (List args))
       end
       else
         push vm (Sx_ref.cek_call f (List args)))
  | Component _ | Island _ ->
    (* Components use keyword-arg parsing — CEK handles this *)
    incr _vm_cek_count;
    let result = Sx_ref.cek_call f (List args) in
    push vm result
  | _ ->
    raise (Eval_error ("VM: not callable: " ^ Sx_runtime.value_to_str f))

(** Main execution loop — iterative (no OCaml stack growth).
    VmClosure calls push frames; the loop picks them up.
    OP_TAIL_CALL + VmClosure = true TCO: drop frame, push new, loop. *)
and run vm =
  while vm.frames <> [] do
    match vm.frames with
    | [] -> ()  (* guard handled by while condition *)
    | frame :: rest_frames ->
      let bc = frame.closure.vm_code.vc_bytecode in
      let consts = frame.closure.vm_code.vc_constants in
      if frame.ip >= Array.length bc then
        vm.frames <- []  (* bytecode exhausted — stop *)
      else begin
        let saved_ip = frame.ip in
        let op = bc.(frame.ip) in
        frame.ip <- frame.ip + 1;
        (try match op with
        (* ---- Constants ---- *)
        | 1 (* OP_CONST *) ->
          let idx = read_u16 frame in
          if idx >= Array.length consts then
            raise (Eval_error (Printf.sprintf "VM: CONST index %d out of bounds (pool size %d)"
              idx (Array.length consts)));
          push vm consts.(idx)
        | 2 (* OP_NIL *)   -> push vm Nil
        | 3 (* OP_TRUE *)  -> push vm (Bool true)
        | 4 (* OP_FALSE *) -> push vm (Bool false)
        | 5 (* OP_POP *)   -> ignore (pop vm)
        | 6 (* OP_DUP *)   -> push vm (peek vm)

        (* ---- Variable access ---- *)
        | 16 (* OP_LOCAL_GET *) ->
          let slot = read_u8 frame in
          let v = match Hashtbl.find_opt frame.local_cells slot with
            | Some cell -> cell.uv_value
            | None ->
              let idx = frame.base + slot in
              if idx >= vm.sp then
                raise (Eval_error (Printf.sprintf
                  "VM: LOCAL_GET slot=%d base=%d sp=%d out of bounds" slot frame.base vm.sp));
              vm.stack.(idx)
          in
          push vm v
        | 17 (* OP_LOCAL_SET *) ->
          let slot = read_u8 frame in
          let v = peek vm in
          (* Write to shared cell if captured, else to stack *)
          (match Hashtbl.find_opt frame.local_cells slot with
           | Some cell -> cell.uv_value <- v
           | None -> vm.stack.(frame.base + slot) <- v)
        | 18 (* OP_UPVALUE_GET *) ->
          let idx = read_u8 frame in
          if idx >= Array.length frame.closure.vm_upvalues then
            raise (Eval_error (Printf.sprintf
              "VM: UPVALUE_GET idx=%d out of bounds (have %d)" idx
              (Array.length frame.closure.vm_upvalues)));
          push vm frame.closure.vm_upvalues.(idx).uv_value
        | 19 (* OP_UPVALUE_SET *) ->
          let idx = read_u8 frame in
          frame.closure.vm_upvalues.(idx).uv_value <- peek vm
        | 20 (* OP_GLOBAL_GET *) ->
          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 ->
            (* 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
        | 21 (* OP_GLOBAL_SET *) ->
          let idx = read_u16 frame in
          let name = match consts.(idx) with String s -> s | _ -> "" in
          (* 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)

        (* ---- Control flow ---- *)
        | 32 (* OP_JUMP *) ->
          let offset = read_i16 frame in
          frame.ip <- frame.ip + offset
        | 33 (* OP_JUMP_IF_FALSE *) ->
          let offset = read_i16 frame in
          let v = pop vm in
          if not (sx_truthy v) then frame.ip <- frame.ip + offset
        | 34 (* OP_JUMP_IF_TRUE *) ->
          let offset = read_i16 frame in
          let v = pop vm in
          if sx_truthy v then frame.ip <- frame.ip + offset

        (* ---- Function calls ---- *)
        | 48 (* OP_CALL *) ->
          let argc = read_u8 frame in
          let args = Array.init argc (fun _ -> pop vm) in
          let f = pop vm in
          let args_list = List.rev (Array.to_list args) in
          vm_call vm f args_list
          (* Loop continues — if VmClosure, new frame runs next iteration *)
        | 49 (* OP_TAIL_CALL *) ->
          let argc = read_u8 frame in
          let args = Array.init argc (fun _ -> pop vm) in
          let f = pop vm in
          let args_list = List.rev (Array.to_list args) in
          (* Drop current frame, reuse stack space — true TCO for VmClosure *)
          vm.frames <- rest_frames;
          vm.sp <- frame.base;
          vm_call vm f args_list
        | 50 (* OP_RETURN *) ->
          let result = pop vm in
          vm.frames <- rest_frames;
          vm.sp <- frame.base;
          push vm result
          (* Loop continues with caller frame *)
        | 51 (* OP_CLOSURE *) ->
          let idx = read_u16 frame in
          if idx >= Array.length consts then
            raise (Eval_error (Printf.sprintf "VM: CLOSURE idx %d >= consts %d" idx (Array.length consts)));
          let code_val = consts.(idx) in
          let code = code_from_value code_val in
          (* Read upvalue descriptors from bytecode *)
          let uv_count = match code_val with
            | Dict d -> (match Hashtbl.find_opt d "upvalue-count" with
                | Some (Number n) -> int_of_float n | _ -> 0)
            | _ -> 0
          in
          let upvalues = Array.init uv_count (fun _ ->
            let is_local = read_u8 frame in
            let index = read_u8 frame in
            if is_local = 1 then begin
              (* Capture from enclosing frame's local slot.
                 Create a shared cell — both parent and closure
                 read/write through this cell. *)
              let cell = match Hashtbl.find_opt frame.local_cells index with
                | Some existing -> existing  (* reuse existing cell *)
                | None ->
                  let c = { uv_value = vm.stack.(frame.base + index) } in
                  Hashtbl.replace frame.local_cells index c;
                  c
              in
              cell
            end else
              (* 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; vm_closure_env = None } in
          push vm (VmClosure cl)
        | 52 (* OP_CALL_PRIM *) ->
          let idx = read_u16 frame in
          let argc = read_u8 frame in
          let name = match consts.(idx) with String s -> s | _ -> "" in
          let args = List.init argc (fun _ -> pop vm) |> List.rev in
          (* Resolve thunks — the CEK evaluator does this automatically
             via trampoline, but the VM must do it explicitly before
             passing args to primitives. *)
          let args = List.map (fun v ->
            match v with
            | Thunk _ -> !Sx_primitives._sx_trampoline_fn v
            | _ -> v) args in
          let result =
            try
              (* Check primitives FIRST (native implementations of map/filter/etc.),
                 then globals (which may have ho_via_cek wrappers that route
                 through the CEK — these can't call VM closures). *)
              let fn_val = try Sx_primitives.get_primitive name with _ ->
                try Hashtbl.find vm.globals name with Not_found ->
                  raise (Eval_error ("VM: unknown primitive " ^ name))
              in
              (match fn_val with
               | NativeFn (_, fn) -> fn args
               | _ -> Nil)
            with Eval_error msg ->
              raise (Eval_error (Printf.sprintf "%s (in CALL_PRIM \"%s\" with %d args)"
                msg name argc))
          in
          push vm result

        (* ---- Collections ---- *)
        | 64 (* OP_LIST *) ->
          let count = read_u16 frame in
          let items = List.init count (fun _ -> pop vm) |> List.rev in
          push vm (List items)
        | 65 (* OP_DICT *) ->
          let count = read_u16 frame in
          let d = Hashtbl.create count in
          for _ = 1 to count do
            let v = pop vm in
            let k = pop vm in
            let key = match k with String s -> s | Keyword s -> s | _ -> Sx_runtime.value_to_str k in
            Hashtbl.replace d key v
          done;
          push vm (Dict d)

        (* ---- String ops ---- *)
        | 144 (* OP_STR_CONCAT *) ->
          let count = read_u8 frame in
          let parts = List.init count (fun _ -> pop vm) |> List.rev in
          let s = String.concat "" (List.map Sx_runtime.value_to_str parts) in
          push vm (String s)

        (* ---- Define ---- *)
        | 128 (* OP_DEFINE *) ->
          let idx = read_u16 frame in
          let name = match consts.(idx) with String s -> s | _ -> "" in
          let v = peek vm in
          Hashtbl.replace vm.globals name v

        (* ---- Inline primitives (no hashtable lookup) ---- *)
        | 160 (* OP_ADD *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Number x, Number y -> Number (x +. y)
            | String x, String y -> String (x ^ y)
            | _ -> Sx_primitives.(get_primitive "+" |> function NativeFn (_, f) -> f [a; b] | _ -> Nil))
        | 161 (* OP_SUB *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with Number x, Number y -> Number (x -. y) | _ -> Nil)
        | 162 (* OP_MUL *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with Number x, Number y -> Number (x *. y) | _ -> Nil)
        | 163 (* OP_DIV *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with Number x, Number y -> Number (x /. y) | _ -> Nil)
        | 164 (* OP_EQ *) ->
          let b = pop vm and a = pop vm in
          push vm (Bool (a = b))
        | 165 (* OP_LT *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with Number x, Number y -> Bool (x < y) | String x, String y -> Bool (x < y) | _ -> Bool false)
        | 166 (* OP_GT *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with Number x, Number y -> Bool (x > y) | String x, String y -> Bool (x > y) | _ -> Bool false)
        | 167 (* OP_NOT *) ->
          let v = pop vm in
          push vm (Bool (not (sx_truthy v)))
        | 168 (* OP_LEN *) ->
          let v = pop vm in
          push vm (match v with
            | List l | ListRef { contents = l } -> Number (float_of_int (List.length l))
            | String s -> Number (float_of_int (String.length s))
            | Dict d -> Number (float_of_int (Hashtbl.length d))
            | Nil -> Number 0.0 | _ -> Number 0.0)
        | 169 (* OP_FIRST *) ->
          let v = pop vm in
          push vm (match v with List (x :: _) | ListRef { contents = x :: _ } -> x | _ -> Nil)
        | 170 (* OP_REST *) ->
          let v = pop vm in
          push vm (match v with List (_ :: xs) | ListRef { contents = _ :: xs } -> List xs | _ -> List [])
        | 171 (* OP_NTH *) ->
          let n = pop vm and coll = pop vm in
          let i = match n with Number f -> int_of_float f | _ -> 0 in
          push vm (match coll with
            | List l | ListRef { contents = l } ->
              (try List.nth l i with _ -> Nil)
            | _ -> Nil)
        | 172 (* OP_CONS *) ->
          let coll = pop vm and x = pop vm in
          push vm (match coll with
            | List l -> List (x :: l)
            | ListRef { contents = l } -> List (x :: l)
            | Nil -> List [x]
            | _ -> List [x])
        | 173 (* OP_NEG *) ->
          let v = pop vm in
          push vm (match v with Number x -> Number (-.x) | _ -> Nil)
        | 174 (* OP_INC *) ->
          let v = pop vm in
          push vm (match v with Number x -> Number (x +. 1.0) | _ -> Nil)
        | 175 (* OP_DEC *) ->
          let v = pop vm in
          push vm (match v with Number x -> Number (x -. 1.0) | _ -> Nil)

        | opcode ->
          raise (Eval_error (Printf.sprintf "VM: unknown opcode %d at ip=%d"
                               opcode (frame.ip - 1)))
        with Invalid_argument msg ->
          let fn_name = match frame.closure.vm_name with Some n -> n | None -> "?" in
          raise (Eval_error (Printf.sprintf
            "VM: %s at ip=%d op=%d in %s (base=%d sp=%d bc_len=%d consts=%d)"
            msg saved_ip op fn_name frame.base vm.sp
            (Array.length bc) (Array.length consts))))
      end
  done

(** 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; 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;
  vm.frames <- [frame];
  run vm;
  pop vm


(** {1 Lazy JIT compilation} *)

(** Compile a lambda or component body to bytecode using the SX compiler.
    Invokes [compile] from spec/compiler.sx via the CEK machine.
    Returns a [vm_closure] ready for execution, or [None] on failure
    (safe fallback to CEK interpretation).

    The compilation cost is a single CEK evaluation of the compiler —
    microseconds per function. The result is cached in the lambda/component
    record so subsequent calls go straight to the VM. *)
let jit_compile_lambda (l : lambda) globals =
  let fn_name = match l.l_name with Some n -> n | None -> "<anon>" in
  try
    let compile_fn = try Hashtbl.find globals "compile"
      with Not_found -> raise (Eval_error "JIT: compiler not loaded") in
    (* Reconstruct the (fn (params) body) form so the compiler produces
       a proper closure. l.l_body is the inner body; we need the full
       function form with params so the compiled code binds them. *)
    let param_syms = List (List.map (fun s -> Symbol s) l.l_params) in
    let fn_expr = List [Symbol "fn"; param_syms; l.l_body] in
    let quoted = List [Symbol "quote"; fn_expr] in
    let result = Sx_ref.eval_expr (List [compile_fn; quoted]) (Env (make_env ())) in
    (* If the lambda has closure-captured variables, merge them into globals
       so the VM can find them via GLOBAL_GET. The compiler doesn't know
       about the enclosing scope, so closure vars get compiled as globals. *)
    let effective_globals =
      let closure = l.l_closure in
      if Hashtbl.length closure.bindings = 0 && closure.parent = None then
        globals  (* no closure vars — use globals directly *)
      else begin
        (* Merge: closure bindings layered on top of globals.
           Use a shallow copy so we don't pollute the real globals. *)
        let merged = Hashtbl.copy globals in
        let rec inject env =
          Hashtbl.iter (fun k v -> Hashtbl.replace merged k v) env.bindings;
          match env.parent with Some p -> inject p | None -> ()
        in
        inject closure;
        let n = Hashtbl.length merged - Hashtbl.length globals in
        if n > 0 then
          Printf.eprintf "[jit] %s: injected %d closure bindings\n%!" fn_name n;
        merged
      end
    in
    (match result with
    | Dict d when Hashtbl.mem d "bytecode" ->
      let outer_code = code_from_value result in
      let bc = outer_code.vc_bytecode in
      if Array.length bc >= 4 && bc.(0) = 51 (* OP_CLOSURE *) then begin
        let idx = bc.(1) lor (bc.(2) lsl 8) in
        if idx < Array.length outer_code.vc_constants then
          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_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);

          None
        end
      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
          let value = execute_module outer_code globals in
          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);
          (* 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
        with _ ->
          Printf.eprintf "[jit] SKIP %s: non-closure execution failed (bc[0]=%d, len=%d)\n%!"
            fn_name (if Array.length bc > 0 then bc.(0) else -1) (Array.length bc);
          None)
      end
    | _ ->
      Printf.eprintf "[jit] FAIL %s: compiler returned %s\n%!" fn_name (type_of result);
      None)
  with e ->
    Printf.eprintf "[jit] FAIL %s: %s\n%!" fn_name (Printexc.to_string e);
    None

(* Wire up forward references *)
let () = jit_compile_ref := jit_compile_lambda
let () = _vm_call_closure_ref := (fun cl args -> call_closure cl args cl.vm_env_ref)