(** 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 *)
}

(** Exception handler entry on the handler stack. *)
type handler_entry = {
  h_catch_ip    : int;    (* IP to jump to when exception is raised *)
  h_frame_depth : int;    (* number of frames when handler was pushed *)
  h_sp          : int;    (* stack pointer when handler was pushed *)
  h_frame       : frame;  (* the frame that pushed the handler *)
}

(** 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 *)
  mutable pending_cek : value option;  (* suspended CEK state from Component/Lambda call *)
  mutable handler_stack : handler_entry list;  (* exception handler stack *)
  mutable reuse_stack : (frame list * int) list;  (* saved call_closure_reuse continuations *)
}

(** Raised when OP_PERFORM is executed. Carries the IO request dict
    and a reference to the VM (which is in a resumable state:
    ip past OP_PERFORM, stack ready for a result push). *)
exception VmSuspended of value * vm

(* Register the VM suspension converter so sx_runtime.sx_apply_cek can
   catch VmSuspended and convert it to CekPerformRequest without a
   direct dependency on this module. *)
let () = Sx_types._convert_vm_suspension := (fun exn ->
  match exn with
  | VmSuspended (request, _vm) -> raise (CekPerformRequest request)
  | _ -> ())


(** 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_rest_arity = -1; vc_locals = 0; vc_bytecode = [||]; vc_constants = [||];
              vc_bytecode_list = None; vc_constants_list = None };
  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

(** Current active VM — allows HO primitives (map, filter, for-each, some)
    to call VmClosure callbacks on the same VM instead of creating a new one.
    This is critical: creating a new VM per callback loses the calling VM's
    stack/frame context, causing upvalue-captured host objects to become
    inaccessible. *)
let _active_vm : vm option ref = ref None

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

(** 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 *)

(** Parse keyword args from an evaluated args list.
    The compiler converts :keyword to its string name, so we need the
    component's param list to identify which strings are keyword names.
    Returns (kwargs_hashtbl, children_list). *)
let parse_keyword_args params args =
  let param_set = Hashtbl.create (List.length params) in
  List.iter (fun p -> Hashtbl.replace param_set p true) params;
  let kwargs = Hashtbl.create 8 in
  let children = ref [] in
  let rec go = function
    | (String k | Keyword k) :: v :: rest when Hashtbl.mem param_set k ->
      Hashtbl.replace kwargs k v; go rest
    | v :: rest -> children := v :: !children; go rest
    | [] -> ()
  in
  go args;
  (kwargs, List.rev !children)

let _vm_comp_jit_count = ref 0
let _vm_comp_cek_count = ref 0
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;
  _vm_comp_jit_count := 0; _vm_comp_cek_count := 0
let vm_report_counters () =
  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

(** 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.
    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
  let rest_arity = cl.vm_code.vc_rest_arity in
  if rest_arity >= 0 then begin
    (* &rest function: push positional args, collect remainder into a list.
       For (fn (a b &rest c) body) with rest_arity=2:
         slots: 0=a, 1=b, 2=c (the rest list) *)
    let nargs = List.length args in
    let rec push_args i = function
      | [] ->
        for _ = i to rest_arity - 1 do push vm Nil done;
        push vm (List [])
      | a :: remaining ->
        if i < rest_arity then (push vm a; push_args (i + 1) remaining)
        else push vm (List (a :: remaining))
    in
    push_args 0 args;
    let used = (if nargs > rest_arity then rest_arity + 1 else nargs + 1) in
    for _ = used to cl.vm_code.vc_locals - 1 do push vm Nil done
  end else begin
    List.iter (fun a -> push vm a) args;
    for _ = List.length args to cl.vm_code.vc_locals - 1 do push vm Nil done
  end;
  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 ->
    (* Accept both compiler output keys (bytecode/constants/arity) and
       SX vm-code keys (vc-bytecode/vc-constants/vc-arity) *)
    let find2 k1 k2 = match Hashtbl.find_opt d k1 with
      | Some _ as r -> r | None -> Hashtbl.find_opt d k2 in
    let bc_list = match find2 "bytecode" "vc-bytecode" with
      | Some (List l | ListRef { contents = l }) ->
        Array.of_list (List.map (fun x -> match x with
          | Integer n -> n | Number n -> int_of_float n | _ -> 0) l)
      | _ -> [||]
    in
    let entries = match find2 "constants" "vc-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" || Hashtbl.mem ed "vc-bytecode" -> entry
      | _ -> entry
    ) entries in
    let arity = match find2 "arity" "vc-arity" with
      | Some (Integer n) -> n | Some (Number n) -> int_of_float n | _ -> 0
    in
    let rest_arity = match find2 "rest-arity" "vc-rest-arity" with
      | Some (Integer n) -> n | Some (Number n) -> int_of_float n | _ -> -1
    in
    (* Compute locals from bytecode: scan for highest LOCAL_GET/LOCAL_SET slot.
       The compiler's arity may undercount when nested lets add many locals. *)
    let max_local = ref (arity - 1) in
    let len = Array.length bc_list in
    let i = ref 0 in
    while !i < len do
      let op = bc_list.(!i) in
      if (op = 16 (* LOCAL_GET *) || op = 17 (* LOCAL_SET *)) && !i + 1 < len then
        (let slot = bc_list.(!i + 1) in
         if slot > !max_local then max_local := slot;
         i := !i + 2)
      else if op = 18 (* UPVALUE_GET *) || op = 19 (* UPVALUE_SET *)
           || op = 8 (* JUMP_IF_FALSE *) || op = 33 (* JUMP_IF_FALSE_u16 *)
           || op = 34 (* JUMP_IF_TRUE *) then
        i := !i + 2
      else if op = 1 (* CONST *) || op = 20 (* GLOBAL_GET *) || op = 21 (* GLOBAL_SET *)
           || op = 32 (* JUMP *) || op = 51 (* CLOSURE *) || op = 52 (* CALL_PRIM *)
           || op = 64 (* MAKE_LIST *) || op = 65 (* MAKE_DICT *) then
        i := !i + 3  (* u16 operand *)
      else
        i := !i + 1
    done;
    let locals = !max_local + 1 + 16 in  (* +16 headroom for temporaries *)
    { vc_arity = arity; vc_rest_arity = rest_arity; vc_locals = locals;
      vc_bytecode = bc_list; vc_constants = constants;
      vc_bytecode_list = None; vc_constants_list = None }
  | _ -> { vc_arity = 0; vc_rest_arity = -1; vc_locals = 16; vc_bytecode = [||]; vc_constants = [||];
           vc_bytecode_list = None; vc_constants_list = None }

(** JIT-compile a component or island body.
    Wraps body as (fn (param1 param2 ... [children]) body) and compiles.
    Returns Some vm_closure on success, None on failure. *)
let jit_compile_comp ~name ~params ~has_children ~body ~closure globals =
  try
    let _compile_fn = try Hashtbl.find globals "compile"
      with Not_found -> raise (Eval_error "JIT: compiler not loaded") in
    let param_names = params @ (if has_children then ["children"] else []) in
    let param_syms = List (List.map (fun s -> Symbol s) param_names) in
    let fn_expr = List [Symbol "fn"; param_syms; body] in
    let quoted = List [Symbol "quote"; fn_expr] 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;
    let result = Sx_ref.eval_expr (List [Symbol "compile"; quoted]) (Env compile_env) 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 = Some name; vm_env_ref = globals;
                 vm_closure_env = Some closure }
        else None
      end else None
    | _ -> None)
  with e ->
    Printf.eprintf "[jit-comp] FAIL %s: %s\n%!" name (Printexc.to_string e);
    None

(** Call an SX value via CEK, detecting suspension instead of erroring.
    Returns the result value, or raises VmSuspended if CEK suspends.
    Saves the suspended CEK state in vm.pending_cek for later resume. *)
let cek_call_or_suspend vm f args =
  incr _vm_cek_count;
  (* Removed debug trace *)
  let a = match args with Nil -> [] | List l -> l | _ -> [args] in
  (* Replace _active_vm with an empty isolation VM so call_closure_reuse
     inside the CEK pushes onto an empty frame stack rather than the caller's.
     Without this, a VmClosure called from within the CEK (e.g. hs-wait)
     merges frames with the caller's VM (e.g. do-repeat), and on resume
     the VM skips the CEK's remaining continuation (wrong mutation order).
     Using Some(isolation) rather than None keeps the call_closure_reuse
     "Some" path which preserves exception identity in js_of_ocaml. *)
  let saved_active = !_active_vm in
  _active_vm := Some (create vm.globals);
  let state = Sx_ref.continue_with_call f (List a) (Env (Sx_types.make_env ())) (List a) (List []) in
  let final = Sx_ref.cek_step_loop state in
  _active_vm := saved_active;
  match Sx_runtime.get_val final (String "phase") with
  | String "io-suspended" ->
    vm.pending_cek <- Some final;
    raise (VmSuspended (Sx_runtime.get_val final (String "request"), vm))
  | _ -> Sx_ref.cek_value final

(** 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 prev_vm = !_active_vm in
  let vm = create globals in
  _active_vm := Some vm;
  push_closure_frame vm cl args;
  (try run vm with e -> _active_vm := prev_vm; raise e);
  _active_vm := prev_vm;
  pop vm

(** 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
    callbacks run on the same VM, avoiding per-call VM allocation overhead. *)
and call_closure_reuse cl args =
  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: save the caller's continuation on the reuse stack.
          DON'T merge frames — that corrupts the frame chain with nested
          closures. On resume, restore_reuse in resume_vm processes these
          in innermost-first order after the callback finishes. *)
       vm.reuse_stack <- (saved_frames, saved_sp) :: vm.reuse_stack;
       raise e
     | e ->
       vm.frames <- saved_frames;
       vm.sp <- saved_sp;
       raise e);
    vm.frames <- saved_frames;
    (* Snapshot/restore sp around the popped result.
       OP_RETURN normally leaves sp = saved_sp + 1, but the bytecode-exhausted
       path (or a callee that returns a closure whose own RETURN leaves extra
       stack residue) can leave sp inconsistent. Read the result at the
       expected slot and reset sp explicitly so the parent frame's
       intermediate values are not corrupted. *)
    let result =
      if vm.sp > saved_sp then vm.stack.(vm.sp - 1)
      else Nil
    in
    vm.sp <- saved_sp;
    result
  | None ->
    call_closure cl args cl.vm_env_ref

(** 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 — push frame on current VM *)
       push_closure_frame vm cl args
     | Some _ ->
       push vm (cek_call_or_suspend vm f (List args))
     | None ->
       if l.l_name <> None
       then begin
         l.l_compiled <- Some jit_failed_sentinel;
         match !jit_compile_ref l vm.globals with
         | Some cl ->
           l.l_compiled <- Some cl;
           push_closure_frame vm cl args
         | None ->
           push vm (cek_call_or_suspend vm f (List args))
       end
       else
         push vm (cek_call_or_suspend vm f (List args)))
  | Component c ->
    let (kwargs, children) = parse_keyword_args c.c_params args in
    (* Get or compile the component body *)
    let compiled = match c.c_compiled with
      | Some cl when not (is_jit_failed cl) -> Some cl
      | Some _ -> None
      | None ->
        c.c_compiled <- Some jit_failed_sentinel;
        let result = jit_compile_comp ~name:c.c_name ~params:c.c_params
          ~has_children:c.c_has_children ~body:c.c_body
          ~closure:c.c_closure vm.globals in
        (match result with Some cl -> c.c_compiled <- Some cl | None -> ());
        result
    in
    (match compiled with
     | Some cl ->
       incr _vm_comp_jit_count;
       (* Build positional args: keyword params in order, then children *)
       let call_args = List.map (fun p ->
         match Hashtbl.find_opt kwargs p with Some v -> v | None -> Nil
       ) c.c_params in
       let call_args = if c.c_has_children
         then call_args @ [List children]
         else call_args in
       (try push vm (call_closure cl call_args cl.vm_env_ref)
        with _ ->
          incr _vm_cek_count; incr _vm_comp_cek_count;
          push vm (cek_call_or_suspend vm f (List args)))
     | None ->
       incr _vm_cek_count; incr _vm_comp_cek_count;
       push vm (cek_call_or_suspend vm f (List args)))
  | Island i ->
    let (kwargs, children) = parse_keyword_args i.i_params args in
    let compiled = match i.i_compiled with
      | Some cl when not (is_jit_failed cl) -> Some cl
      | Some _ -> None
      | None ->
        i.i_compiled <- Some jit_failed_sentinel;
        let result = jit_compile_comp ~name:i.i_name ~params:i.i_params
          ~has_children:i.i_has_children ~body:i.i_body
          ~closure:i.i_closure vm.globals in
        (match result with Some cl -> i.i_compiled <- Some cl | None -> ());
        result
    in
    (match compiled with
     | Some cl ->
       incr _vm_comp_jit_count;
       let call_args = List.map (fun p ->
         match Hashtbl.find_opt kwargs p with Some v -> v | None -> Nil
       ) i.i_params in
       let call_args = if i.i_has_children
         then call_args @ [List children]
         else call_args in
       (try push vm (call_closure cl call_args cl.vm_env_ref)
        with _ ->
          incr _vm_cek_count; incr _vm_comp_cek_count;
          push vm (cek_call_or_suspend vm f (List args)))
     | None ->
       incr _vm_cek_count; incr _vm_comp_cek_count;
       push vm (cek_call_or_suspend vm f (List args)))
  | _ ->
    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 begin
        (* Bytecode exhausted without explicit RETURN — pop frame like RETURN *)
        let fn_name = match frame.closure.vm_name with Some n -> n | None -> "?" in
        Printf.eprintf "[vm] WARN: bytecode exhausted without RETURN in %s (base=%d sp=%d frames=%d)\n%!"
          fn_name frame.base vm.sp (List.length rest_frames);
        let result = if vm.sp > frame.base then pop vm else Nil in
        vm.frames <- rest_frames;
        vm.sp <- frame.base;
        if rest_frames <> [] then push vm result
        (* If no more frames, result stays on stack for call_closure to pop *)
      end
      else begin
        let saved_ip = frame.ip in
        let op = bc.(frame.ip) in
        frame.ip <- frame.ip + 1;
        incr _vm_insn_count;
        (* Check timeout — compare VM instruction count against step limit *)
        if !_vm_insn_count land 0xFFFF = 0 && !Sx_ref.step_limit > 0
           && !_vm_insn_count > !Sx_ref.step_limit then
          raise (Eval_error "TIMEOUT: step limit exceeded");
        (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)
        | 7 (* OP_SWAP *)  ->
          let a = pop vm in let b = pop vm in
          push vm a; push vm b

        (* ---- 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
          (* Check closure env first (matches OP_GLOBAL_SET priority) *)
          let id = Sx_types.intern name in
          let found_in_env = match frame.closure.vm_closure_env with
            | Some env ->
              let rec env_lookup e =
                try Some (Hashtbl.find e.bindings id)
                with Not_found ->
                  match e.parent with Some p -> env_lookup p | None -> None
              in env_lookup env
            | None -> None
          in
          let v = match found_in_env with
            | Some v -> v
            | None ->
              try Hashtbl.find vm.globals name with Not_found ->
                try Sx_primitives.get_primitive name
                with _ ->
                  (* Try resolve hook — loads the library that exports this symbol *)
                  (try
                    let resolve_fn = Hashtbl.find vm.globals "__resolve-symbol" in
                    ignore (Sx_runtime.sx_call resolve_fn [String name]);
                    try Hashtbl.find vm.globals name
                    with Not_found -> raise (Eval_error ("VM undefined: " ^ name))
                  with Not_found -> 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 id = Sx_types.intern name in
              let rec find_env e =
                if Hashtbl.mem e.bindings id then
                  (Hashtbl.replace e.bindings id (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 begin
            let v = peek vm in
            Hashtbl.replace vm.globals name v;
            (match !Sx_types._vm_global_set_hook with Some f -> f name v | None -> ())
          end

        (* ---- 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

        (* ---- Exception handling ---- *)
        | 35 (* OP_PUSH_HANDLER *) ->
          let catch_offset = read_i16 frame in
          let entry = {
            h_catch_ip = frame.ip + catch_offset;
            h_frame_depth = List.length vm.frames;
            h_sp = vm.sp;
            h_frame = frame;
          } in
          vm.handler_stack <- entry :: vm.handler_stack
        | 36 (* OP_POP_HANDLER *) ->
          (match vm.handler_stack with
           | _ :: rest -> vm.handler_stack <- rest
           | [] -> ())
        | 37 (* OP_RAISE *) ->
          let exn_val = pop vm in
          (match vm.handler_stack with
           | entry :: rest ->
             vm.handler_stack <- rest;
             (* Unwind frames to the handler's depth *)
             while List.length vm.frames > entry.h_frame_depth do
               match vm.frames with
               | _ :: fs -> vm.frames <- fs
               | [] -> ()
             done;
             (* Restore stack pointer and jump to catch *)
             vm.sp <- entry.h_sp;
             entry.h_frame.ip <- entry.h_catch_ip;
             push vm exn_val
           | [] ->
             (* No handler — raise OCaml exception for CEK to catch *)
             raise (Eval_error (Printf.sprintf "Unhandled exception: %s"
               (Sx_runtime.value_to_str exn_val))))

        (* ---- 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 = frame.closure.vm_closure_env } 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
              (* Single lookup: vm.globals is the sole source of truth.
                 Primitives are seeded into vm.globals at init as NativeFn values.
                 OP_DEFINE and registerNative naturally override them. *)
              let fn_val = try Hashtbl.find vm.globals name with Not_found ->
                (* Fallback to Sx_primitives — primitives registered AFTER JIT
                   setup (e.g. host-global, host-get registered inside the test
                   runner's bind/register path) are not in vm.globals. *)
                try Sx_primitives.get_primitive name
                with _ ->
                  raise (Eval_error ("VM: unknown primitive " ^ name))
              in
              (match fn_val with
               | NativeFn (_, fn) -> fn args
               | VmClosure _ | Lambda _ | Component _ | Island _ ->
                 Sx_ref.cek_call fn_val (List 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;
          (match !Sx_types._vm_global_set_hook with
           | Some f -> f name v | None -> ())

        (* ---- Inline primitives ----
           Fast path for common types; fallback to actual primitive
           for edge cases (type coercion, thunks, RawHTML, etc.)
           to guarantee behavioral parity with CALL_PRIM. *)
        | 160 (* OP_ADD *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Integer x, Integer y -> Integer (x + y)
            | Number x, Number y -> Number (x +. y)
            | Integer x, Number y -> Number (float_of_int x +. y)
            | Number x, Integer y -> Number (x +. float_of_int y)
            | _ -> (Hashtbl.find Sx_primitives.primitives "+") [a; b])
        | 161 (* OP_SUB *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Integer x, Integer y -> Integer (x - y)
            | Number x, Number y -> Number (x -. y)
            | Integer x, Number y -> Number (float_of_int x -. y)
            | Number x, Integer y -> Number (x -. float_of_int y)
            | _ -> (Hashtbl.find Sx_primitives.primitives "-") [a; b])
        | 162 (* OP_MUL *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Integer x, Integer y -> Integer (x * y)
            | Number x, Number y -> Number (x *. y)
            | Integer x, Number y -> Number (float_of_int x *. y)
            | Number x, Integer y -> Number (x *. float_of_int y)
            | _ -> (Hashtbl.find Sx_primitives.primitives "*") [a; b])
        | 163 (* OP_DIV *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Integer x, Integer y when y <> 0 && x mod y = 0 -> Integer (x / y)
            | Integer x, Integer y -> Number (float_of_int x /. float_of_int y)
            | Number x, Number y -> Number (x /. y)
            | Integer x, Number y -> Number (float_of_int x /. y)
            | Number x, Integer y -> Number (x /. float_of_int y)
            | _ -> (Hashtbl.find Sx_primitives.primitives "/") [a; b])
        | 164 (* OP_EQ *) ->
          let b = pop vm and a = pop vm in
          push vm (Bool (Sx_runtime._fast_eq a b))
        | 165 (* OP_LT *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Integer x, Integer y -> Bool (x < y)
            | Number x, Number y -> Bool (x < y)
            | Integer x, Number y -> Bool (float_of_int x < y)
            | Number x, Integer y -> Bool (x < float_of_int y)
            | String x, String y -> Bool (x < y)
            | _ -> Sx_runtime.prim_call "<" [a; b])
        | 166 (* OP_GT *) ->
          let b = pop vm and a = pop vm in
          push vm (match a, b with
            | Integer x, Integer y -> Bool (x > y)
            | Number x, Number y -> Bool (x > y)
            | Integer x, Number y -> Bool (float_of_int x > y)
            | Number x, Integer y -> Bool (x > float_of_int y)
            | String x, String y -> Bool (x > y)
            | _ -> Sx_runtime.prim_call ">" [a; b])
        | 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 } -> Integer (List.length l)
            | String s -> Integer (String.length s)
            | Dict d -> Integer (Hashtbl.length d)
            | Nil -> Integer 0
            | _ -> (Hashtbl.find Sx_primitives.primitives "len") [v])
        | 169 (* OP_FIRST *) ->
          let v = pop vm in
          push vm (match v with
            | List (x :: _) | ListRef { contents = x :: _ } -> x
            | List [] | ListRef { contents = [] } | Nil -> Nil
            | _ -> (Hashtbl.find Sx_primitives.primitives "first") [v])
        | 170 (* OP_REST *) ->
          let v = pop vm in
          push vm (match v with
            | List (_ :: xs) | ListRef { contents = _ :: xs } -> List xs
            | List [] | ListRef { contents = [] } | Nil -> List []
            | _ -> (Hashtbl.find Sx_primitives.primitives "rest") [v])
        | 171 (* OP_NTH *) ->
          let n = pop vm and coll = pop vm in
          push vm (match coll, n with
            | (List l | ListRef { contents = l }), Number f ->
              (try List.nth l (int_of_float f) with _ -> Nil)
            | String s, Number f ->
              let i = int_of_float f in
              if i >= 0 && i < String.length s then String (String.make 1 s.[i])
              else Nil
            | _ -> (Hashtbl.find Sx_primitives.primitives "nth") [coll; n])
        | 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]
            | _ -> (Hashtbl.find Sx_primitives.primitives "cons") [x; coll])
        | 173 (* OP_NEG *) ->
          let v = pop vm in
          push vm (match v with
            | Number x -> Number (-.x)
            | _ -> (Hashtbl.find Sx_primitives.primitives "-") [v])
        | 174 (* OP_INC *) ->
          let v = pop vm in
          push vm (match v with
            | Number x -> Number (x +. 1.0)
            | _ -> (Hashtbl.find Sx_primitives.primitives "inc") [v])
        | 175 (* OP_DEC *) ->
          let v = pop vm in
          push vm (match v with
            | Number x -> Number (x -. 1.0)
            | _ -> (Hashtbl.find Sx_primitives.primitives "dec") [v])

        (* ---- IO Suspension ---- *)
        | 112 (* OP_PERFORM *) ->
          let request = pop vm in
          raise (VmSuspended (request, vm))

        | 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

(** Resume a suspended VM by pushing the IO result and continuing.
    May raise VmSuspended again if the VM hits another OP_PERFORM.

    After the callback finishes, restores any call_closure_reuse
    continuations saved on vm.reuse_stack (innermost first). *)
let resume_vm vm result =
  (match vm.pending_cek with
   | Some cek_state ->
     vm.pending_cek <- None;
     let final = Sx_ref.cek_resume cek_state result in
     (match Sx_runtime.get_val final (String "phase") with
      | String "io-suspended" ->
        vm.pending_cek <- Some final;
        raise (VmSuspended (Sx_runtime.get_val final (String "request"), vm))
      | _ ->
        push vm (Sx_ref.cek_value final))
   | None ->
     push vm result);
  (try run vm
   with
   | VmSuspended _ as e ->
     (* Re-suspension during resume: the VM hit another perform. *)
     raise e
   | Eval_error msg ->
     (* Error during resumed execution. If the VM has a handler on its
        handler_stack, dispatch to it (same as OP_RAISE). This enables
        try/catch across async perform/resume boundaries — the handler
        was pushed before the perform and survives on the vm struct. *)
     (match vm.handler_stack with
      | entry :: rest ->
        vm.handler_stack <- rest;
        while List.length vm.frames > entry.h_frame_depth do
          match vm.frames with _ :: fs -> vm.frames <- fs | [] -> ()
        done;
        vm.sp <- entry.h_sp;
        entry.h_frame.ip <- entry.h_catch_ip;
        push vm (String msg);
        run vm
      | [] -> raise (Eval_error msg)));
  (* Clear reuse_stack — any entries here are stale from the original
     suspension and don't apply to the current state. The VM just
     completed its execution successfully. *)
  vm.reuse_stack <- [];
  (* Restore call_closure_reuse continuations saved during suspension.
     reuse_stack is in catch order (outermost first from prepend) —
     reverse to get innermost first, matching callback→caller unwinding. *)
  let rec restore_reuse pending =
    match pending with
    | [] -> ()
    | (saved_frames, saved_sp) :: rest ->
      let callback_result = pop vm in
      vm.frames <- saved_frames;
      (* Restore sp to the value captured before the suspended callee was
         pushed. The callee's locals/temps may still be on the stack above
         saved_sp; without this reset, subsequent LOCAL_GET/SET in the
         caller frame (e.g. letrec sibling bindings waiting on the call)
         see stale callee data instead of their own slots. Mirrors the
         OP_RETURN+sp-reset semantics that sync `call_closure_reuse`
         relies on for clean caller-frame state. *)
      if saved_sp < vm.sp then vm.sp <- saved_sp;
      push vm callback_result;
      (try
        run vm;
        (* Check for new reuse entries added by nested call_closure_reuse *)
        let new_pending = List.rev vm.reuse_stack in
        vm.reuse_stack <- [];
        restore_reuse (new_pending @ rest)
      with VmSuspended _ as e ->
        (* Re-suspension: save unprocessed entries back for next resume.
           rest is innermost-first; vm.reuse_stack is outermost-first.
           Combine so next resume's reversal yields: new_inner, old_inner→outer. *)
        vm.reuse_stack <- (List.rev rest) @ vm.reuse_stack;
        raise e)
  in
  let pending = List.rev vm.reuse_stack in
  vm.reuse_stack <- [];
  restore_reuse pending;
  pop vm

(** 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

(** Execute module, catching VmSuspended locally (same compilation unit).
    Returns [Ok result] or [Error (request, vm)] for import suspension.
    Needed because js_of_ocaml can't catch exceptions across module boundaries. *)
let execute_module_safe code globals =
  try
    let result = execute_module code globals in
    Ok result
  with VmSuspended (request, vm) ->
    Error (request, 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. *)
(* Functions whose JIT bytecode is known broken (see project_jit_bytecode_bug):
   parser combinators drop intermediate results, the hyperscript parse/compile
   stack corrupts ASTs when compiled, and test-orchestration helpers have
   call-count/arg-shape mismatches vs CEK. These must run under CEK. *)
let _jit_is_broken_name n =
  (* Parser combinators *)
  n = "parse-bind" || n = "seq" || n = "seq2" || n = "many" || n = "many1"
  || n = "satisfy" || n = "fmap" || n = "alt" || n = "alt2"
  || n = "skip-left" || n = "skip-right" || n = "skip-many" || n = "optional"
  || n = "between" || n = "sep-by" || n = "sep-by1" || n = "parse-char"
  || n = "parse-string" || n = "lazy-parser" || n = "label"
  || n = "not-followed-by" || n = "look-ahead"
  (* Hyperscript orchestrators — call parser combinators *)
  || n = "hs-tokenize" || n = "hs-parse" || n = "hs-compile"
  || n = "hs-to-sx" || n = "hs-to-sx-from-source"
  (* Test orchestration helpers *)
  || n = "eval-hs" || n = "eval-hs-inner" || n = "eval-hs-with-me"
  || n = "run-hs-fixture"
  (* Large top-level functions whose JIT compile exceeds the 5s test
     deadline — tw-resolve-style, tw-resolve-layout, graphql parse. *)
  || n = "tw-resolve-style" || n = "tw-resolve-layout"
  || n = "gql-ws?" || n = "gql-parse-tokens" || n = "gql-execute-operation"
  (* Hyperscript loop runtime: uses `guard` to catch hs-break/hs-continue
     exceptions. JIT-compiled guard drops the exception handler such that
     break propagates out of the click handler instead of exiting the loop.
     See hs-upstream-repeat/hs-upstream-put tests. *)
  || n = "hs-repeat-times" || n = "hs-repeat-forever"
  || n = "hs-repeat-while" || n = "hs-repeat-until"
  || n = "hs-for-each" || n = "hs-put!"

let jit_compile_lambda (l : lambda) globals =
  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 if List.mem "&key" l.l_params || List.mem ":as" l.l_params then (
    (* &key/:as require complex runtime argument processing that the compiler
       doesn't emit. These functions must run via CEK. *)
    None
  ) else if l.l_name = None || l.l_closure.Sx_types.parent <> None then (
    (* Anonymous or nested lambdas: skip JIT. Nested defines get re-created
       on each outer call, so per-call compile cost is pure overhead. *)
    None
  ) else if _jit_is_broken_name fn_name then (
    None
  ) else
  try
    _jit_compiling := true;
    let compile_fn = try Hashtbl.find globals "compile"
      with Not_found -> (_jit_compiling := false; raise (Eval_error "JIT: compiler not loaded")) in
    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
    (* Fast path: if compile has bytecode, call it directly via the VM.
       All helper calls (compile-expr, emit-byte, etc.) happen inside the
       same VM execution — no per-call VM allocation overhead. *)
    let result = match compile_fn with
    | Lambda { l_compiled = Some cl; _ } when not (is_jit_failed cl) ->
      call_closure cl [fn_expr] globals
    | _ ->
      ignore compile_fn;
      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;
      Sx_ref.eval_expr (List [Symbol "compile"; quoted]) (Env compile_env)
    in
    _jit_compiling := false;
    (* Merge closure bindings into effective_globals so GLOBAL_GET resolves
       variables from let/define blocks. The compiler emits GLOBAL_GET for
       free variables; the VM resolves them from vm_env_ref. *)
    let effective_globals =
      if Hashtbl.length l.l_closure.Sx_types.bindings > 0 then begin
        let merged = Hashtbl.copy globals in
        let rec merge_env env =
          Hashtbl.iter (fun id v ->
            let name = Sx_types.unintern id in
            if not (Hashtbl.mem merged name) then
              Hashtbl.replace merged name v) env.Sx_types.bindings;
          match env.Sx_types.parent with Some p -> merge_env p | None -> ()
        in
        merge_env l.l_closure;
        merged
      end else globals
    in
    (match result with
    | Dict d when Hashtbl.mem d "bytecode" || Hashtbl.mem d "vc-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
        (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);
          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 ->
    _jit_compiling := false;
    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_reuse cl args)
let () = _vm_suspension_to_dict := (fun exn ->
  match exn with
  | VmSuspended (request, vm) ->
    (* Snapshot pending_cek and reuse_stack NOW — a nested cek_call_or_suspend
       on the same VM may overwrite them before our resume function is called. *)
    let saved_cek = vm.pending_cek in
    let saved_reuse = vm.reuse_stack in
    let d = Hashtbl.create 3 in
    Hashtbl.replace d "__vm_suspended" (Bool true);
    Hashtbl.replace d "request" request;
    Hashtbl.replace d "resume" (NativeFn ("vm-resume", fun args ->
      match args with
      | [result] ->
        (* Restore saved state before resuming — may have been overwritten
           by a nested suspension on the same VM. *)
        vm.pending_cek <- saved_cek;
        vm.reuse_stack <- saved_reuse;
        (try resume_vm vm result
         with exn2 ->
           match !_vm_suspension_to_dict exn2 with
           | Some marker -> marker
           | None -> raise exn2)
      | _ -> Nil));
    Some (Dict d)
  | _ -> None)
(* Hook: when eval_expr (cek_run_iterative) encounters a CEK suspension,
   convert it to VmSuspended so it propagates to the outer handler
   (value_to_js wrapper, _driveAsync, etc.). Without this, perform
   inside nested eval_expr calls (event handler → trampoline → eval_expr)
   gets swallowed as "IO suspension in non-IO context". *)
let () = _cek_io_suspend_hook := Some (fun suspended_state ->
  let request = Sx_ref.cek_io_request suspended_state in
  let vm = create !_default_vm_globals in
  vm.pending_cek <- Some suspended_state;
  (* Transfer reuse_stack from the active VM so resume_vm can restore
     caller frames saved by call_closure_reuse during the suspension chain. *)
  (match !_active_vm with
   | Some active when active.reuse_stack <> [] ->
     vm.reuse_stack <- active.reuse_stack;
     active.reuse_stack <- []
   | _ -> ());
  raise (VmSuspended (request, vm)))

let () = _cek_eval_lambda_ref := (fun f args ->
  let state = Sx_ref.continue_with_call f (List args) (Env (make_env ())) (List args) (List []) in
  let final = Sx_ref.cek_step_loop state in
  match Sx_runtime.get_val final (String "phase") with
  | String "io-suspended" ->
    (* Create a stub VM to carry the suspended CEK state.
       resume_vm will: cek_resume → push result → run (no-op, no frames) → pop *)
    let vm = create (Hashtbl.create 0) in
    vm.pending_cek <- Some final;
    (* Transfer reuse_stack from active VM *)
    (match !_active_vm with
     | Some active when active.reuse_stack <> [] ->
       vm.reuse_stack <- active.reuse_stack;
       active.reuse_stack <- []
     | _ -> ());
    raise (VmSuspended (Sx_runtime.get_val final (String "request"), vm))
  | _ -> Sx_ref.cek_value final)


(** {1 Debugging / introspection} *)

(** Map opcode integer to human-readable name. *)
let opcode_name = function
  | 1 -> "CONST" | 2 -> "NIL" | 3 -> "TRUE" | 4 -> "FALSE"
  | 5 -> "POP" | 6 -> "DUP" | 7 -> "SWAP"
  | 16 -> "LOCAL_GET" | 17 -> "LOCAL_SET"
  | 18 -> "UPVALUE_GET" | 19 -> "UPVALUE_SET"
  | 20 -> "GLOBAL_GET" | 21 -> "GLOBAL_SET"
  | 32 -> "JUMP" | 33 -> "JUMP_IF_FALSE" | 34 -> "JUMP_IF_TRUE"
  | 35 -> "PUSH_HANDLER" | 36 -> "POP_HANDLER" | 37 -> "RAISE"
  | 48 -> "CALL" | 49 -> "TAIL_CALL" | 50 -> "RETURN"
  | 51 -> "CLOSURE" | 52 -> "CALL_PRIM"
  | 64 -> "LIST" | 65 -> "DICT"
  | 128 -> "DEFINE"
  | 144 -> "STR_CONCAT"
  | 160 -> "ADD" | 161 -> "SUB" | 162 -> "MUL" | 163 -> "DIV"
  | 164 -> "EQ" | 165 -> "LT" | 166 -> "GT" | 167 -> "NOT"
  | 168 -> "LEN" | 169 -> "FIRST" | 170 -> "REST" | 171 -> "NTH"
  | 172 -> "CONS" | 173 -> "NEG" | 174 -> "INC" | 175 -> "DEC"
  | n -> Printf.sprintf "UNKNOWN_%d" n

(** Number of extra operand bytes consumed by each opcode.
    Returns (format, total_bytes) where format describes the operand types. *)
let opcode_operand_size = function
  | 1 (* CONST *) | 20 (* GLOBAL_GET *) | 21 (* GLOBAL_SET *)
  | 64 (* LIST *) | 65 (* DICT *) | 128 (* DEFINE *) -> 2  (* u16 *)
  | 16 (* LOCAL_GET *) | 17 (* LOCAL_SET *)
  | 18 (* UPVALUE_GET *) | 19 (* UPVALUE_SET *)
  | 48 (* CALL *) | 49 (* TAIL_CALL *)
  | 144 (* STR_CONCAT *) -> 1  (* u8 *)
  | 32 (* JUMP *) | 33 (* JUMP_IF_FALSE *) | 34 (* JUMP_IF_TRUE *)
  | 35 (* PUSH_HANDLER *) -> 2  (* i16 *)
  | 51 (* CLOSURE *) -> 2  (* u16 for constant index; upvalue descriptors follow dynamically *)
  | 52 (* CALL_PRIM *) -> 3  (* u16 + u8 *)
  | _ -> 0  (* no operand *)

(** Trace a single execution — compile + run, collecting trace entries.
    Each entry is a dict with :opcode, :stack, :depth. *)
let trace_run src globals =
  (* Compile *)
  let compile_fn = try Hashtbl.find globals "compile"
    with Not_found -> raise (Eval_error "trace: compiler not loaded") in
  let exprs = Sx_parser.parse_all src in
  let expr = match exprs with [e] -> e | _ -> List (Symbol "do" :: exprs) in
  let quoted = List [Symbol "quote"; expr] in
  let code_val = Sx_ref.eval_expr (List [compile_fn; quoted]) (Env (make_env ())) in
  let code = code_from_value code_val in
  let cl = { vm_code = code; vm_upvalues = [||]; vm_name = Some "trace";
             vm_env_ref = globals; vm_closure_env = None } in
  let vm = create globals in
  let frame0 = { 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 <- [frame0];
  (* Run with tracing *)
  let trace = ref [] in
  let max_steps = 10000 in
  let steps = ref 0 in
  (try
    while vm.frames <> [] && !steps < max_steps do
      match vm.frames with
      | [] -> ()
      | frame :: _ ->
        let bc = frame.closure.vm_code.vc_bytecode in
        if frame.ip >= Array.length bc then
          vm.frames <- []
        else begin
          let op = bc.(frame.ip) in
          (* Snapshot stack top 5 *)
          let stack_snap = List.init (min 5 vm.sp) (fun i ->
            let v = vm.stack.(vm.sp - 1 - i) in
            String (Sx_types.inspect v)) in
          let entry = Hashtbl.create 4 in
          Hashtbl.replace entry "opcode" (String (opcode_name op));
          Hashtbl.replace entry "stack" (List stack_snap);
          Hashtbl.replace entry "depth" (Number (float_of_int (List.length vm.frames)));
          trace := Dict entry :: !trace;
          incr steps;
          (* Execute one step — use the main run loop for 1 step.
             We do this by saving the state and running the original dispatch. *)
          let saved_ip = frame.ip in
          frame.ip <- frame.ip + 1;
          let rest_frames = List.tl vm.frames in
          (try match op with
          | 1 -> let idx = read_u16 frame in push vm frame.closure.vm_code.vc_constants.(idx)
          | 2 -> push vm Nil
          | 3 -> push vm (Bool true)
          | 4 -> push vm (Bool false)
          | 5 -> ignore (pop vm)
          | 6 -> push vm (peek vm)
          | 7 -> let a = pop vm in let b = pop vm in push vm a; push vm b
          | 16 -> let slot = read_u8 frame in
            let v = match Hashtbl.find_opt frame.local_cells slot with
              | Some cell -> cell.uv_value
              | None -> vm.stack.(frame.base + slot) in
            push vm v
          | 17 -> let slot = read_u8 frame in let v = peek vm in
            (match Hashtbl.find_opt frame.local_cells slot with
             | Some cell -> cell.uv_value <- v
             | None -> vm.stack.(frame.base + slot) <- v)
          | 18 -> let idx = read_u8 frame in
            push vm frame.closure.vm_upvalues.(idx).uv_value
          | 19 -> let idx = read_u8 frame in
            frame.closure.vm_upvalues.(idx).uv_value <- peek vm
          | 20 -> let idx = read_u16 frame in
            let name = match frame.closure.vm_code.vc_constants.(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)) in
            push vm v
          | 21 -> let idx = read_u16 frame in
            let name = match frame.closure.vm_code.vc_constants.(idx) with String s -> s | _ -> "" in
            Hashtbl.replace vm.globals name (peek vm)
          | 32 -> let offset = read_i16 frame in frame.ip <- frame.ip + offset
          | 33 -> let offset = read_i16 frame in let v = pop vm in
            if not (sx_truthy v) then frame.ip <- frame.ip + offset
          | 34 -> let offset = read_i16 frame in let v = pop vm in
            if sx_truthy v then frame.ip <- frame.ip + offset
          | 35 -> let catch_offset = read_i16 frame in
            vm.handler_stack <- { h_catch_ip = frame.ip + catch_offset;
              h_frame_depth = List.length vm.frames; h_sp = vm.sp;
              h_frame = frame } :: vm.handler_stack
          | 36 -> (match vm.handler_stack with _ :: r -> vm.handler_stack <- r | [] -> ())
          | 37 -> let exn_val = pop vm in
            (match vm.handler_stack with
             | entry :: rest ->
               vm.handler_stack <- rest;
               while List.length vm.frames > entry.h_frame_depth do
                 match vm.frames with _ :: fs -> vm.frames <- fs | [] -> () done;
               vm.sp <- entry.h_sp; entry.h_frame.ip <- entry.h_catch_ip;
               push vm exn_val
             | [] -> vm.frames <- [])
          | 48 -> let argc = read_u8 frame in
            let args = Array.init argc (fun _ -> pop vm) in
            let f = pop vm in
            vm_call vm f (List.rev (Array.to_list args))
          | 49 -> let argc = read_u8 frame in
            let args = Array.init argc (fun _ -> pop vm) in
            let f = pop vm in
            vm.frames <- rest_frames; vm.sp <- frame.base;
            vm_call vm f (List.rev (Array.to_list args))
          | 50 -> let result = pop vm in
            vm.frames <- rest_frames; vm.sp <- frame.base; push vm result
          | 51 -> (* CLOSURE — skip for trace, just advance past upvalue descriptors *)
            let idx = read_u16 frame in
            let code_val2 = frame.closure.vm_code.vc_constants.(idx) in
            let uv_count = match code_val2 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
                let cell = match Hashtbl.find_opt frame.local_cells index with
                  | Some existing -> existing
                  | None ->
                    let c = { uv_value = vm.stack.(frame.base + index) } in
                    Hashtbl.replace frame.local_cells index c; c in
                cell
              end else frame.closure.vm_upvalues.(index)
            ) in
            let inner_code = code_from_value code_val2 in
            let c = { vm_code = inner_code; vm_upvalues = upvalues; vm_name = None;
                       vm_env_ref = vm.globals; vm_closure_env = frame.closure.vm_closure_env } in
            push vm (VmClosure c)
          | 52 -> let idx = read_u16 frame in let argc = read_u8 frame in
            let name = match frame.closure.vm_code.vc_constants.(idx) with String s -> s | _ -> "" in
            let args = List.init argc (fun _ -> pop vm) |> List.rev in
            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) -> push vm (fn args) | _ -> push vm Nil)
          | 64 -> let count = read_u16 frame in
            let items = List.init count (fun _ -> pop vm) |> List.rev in
            push vm (List items)
          | 65 -> 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)
          | 128 -> let idx = read_u16 frame in
            let name = match frame.closure.vm_code.vc_constants.(idx) with String s -> s | _ -> "" in
            Hashtbl.replace vm.globals name (peek vm)
          | 144 -> let count = read_u8 frame in
            let parts = List.init count (fun _ -> pop vm) |> List.rev in
            push vm (String (String.concat "" (List.map Sx_runtime.value_to_str parts)))
          | 160 -> let b = pop vm and a = pop vm in
            push vm (match a, b with Number x, Number y -> Number (x +. y) | _ -> Nil)
          | 161 -> let b = pop vm and a = pop vm in
            push vm (match a, b with Number x, Number y -> Number (x -. y) | _ -> Nil)
          | 162 -> let b = pop vm and a = pop vm in
            push vm (match a, b with Number x, Number y -> Number (x *. y) | _ -> Nil)
          | 163 -> let b = pop vm and a = pop vm in
            push vm (match a, b with Number x, Number y -> Number (x /. y) | _ -> Nil)
          | 164 -> let b = pop vm and a = pop vm in push vm (Bool (a = b))
          | 165 -> let b = pop vm and a = pop vm in
            push vm (match a, b with Number x, Number y -> Bool (x < y) | _ -> Bool false)
          | 166 -> let b = pop vm and a = pop vm in
            push vm (match a, b with Number x, Number y -> Bool (x > y) | _ -> Bool false)
          | 167 -> let v = pop vm in push vm (Bool (not (sx_truthy v)))
          | 168 -> 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))
              | _ -> Number 0.0)
          | 169 -> let v = pop vm in
            push vm (match v with List (x :: _) | ListRef { contents = x :: _ } -> x | _ -> Nil)
          | 170 -> let v = pop vm in
            push vm (match v with
              | List (_ :: xs) | ListRef { contents = _ :: xs } -> List xs | _ -> List [])
          | 171 -> let n = pop vm and coll = pop vm in
            push vm (match coll, n with
              | (List l | ListRef { contents = l }), Number f ->
                (try List.nth l (int_of_float f) with _ -> Nil) | _ -> Nil)
          | 172 -> let coll = pop vm and x = pop vm in
            push vm (match coll with List l -> List (x :: l) | _ -> List [x])
          | 173 -> let v = pop vm in
            push vm (match v with Number x -> Number (-.x) | _ -> Nil)
          | 174 -> let v = pop vm in
            push vm (match v with Number x -> Number (x +. 1.0) | _ -> Nil)
          | 175 -> let v = pop vm in
            push vm (match v with Number x -> Number (x -. 1.0) | _ -> Nil)
          | _ -> ()
          with e ->
            let _ = e in
            ignore saved_ip;
            (* On error during trace, just stop *)
            vm.frames <- [])
        end
    done
  with _ -> ());
  List (List.rev !trace)

(** Disassemble a vm_code into a list of instruction dicts. *)
let disassemble (code : vm_code) =
  let bc = code.vc_bytecode in
  let len = Array.length bc in
  let consts = code.vc_constants in
  let instrs = ref [] in
  let ip = ref 0 in
  while !ip < len do
    let offset = !ip in
    let op = bc.(!ip) in
    ip := !ip + 1;
    let name = opcode_name op in
    let operands = ref [] in
    (match op with
    | 1 (* CONST *) | 20 (* GLOBAL_GET *) | 21 (* GLOBAL_SET *)
    | 128 (* DEFINE *) ->
      if !ip + 1 < len then begin
        let lo = bc.(!ip) in let hi = bc.(!ip + 1) in
        let idx = lo lor (hi lsl 8) in
        ip := !ip + 2;
        let const_str = if idx < Array.length consts
          then Sx_types.inspect consts.(idx) else "?" in
        operands := [Number (float_of_int idx); String const_str]
      end
    | 64 (* LIST *) | 65 (* DICT *) | 51 (* CLOSURE *) ->
      if !ip + 1 < len then begin
        let lo = bc.(!ip) in let hi = bc.(!ip + 1) in
        let idx = lo lor (hi lsl 8) in
        ip := !ip + 2;
        operands := [Number (float_of_int idx)];
        (* For CLOSURE, skip upvalue descriptors *)
        if op = 51 && idx < Array.length consts then begin
          let uv_count = match consts.(idx) with
            | Dict d -> (match Hashtbl.find_opt d "upvalue-count" with
                | Some (Number n) -> int_of_float n | _ -> 0)
            | _ -> 0 in
          ip := !ip + uv_count * 2
        end
      end
    | 16 (* LOCAL_GET *) | 17 (* LOCAL_SET *)
    | 18 (* UPVALUE_GET *) | 19 (* UPVALUE_SET *)
    | 48 (* CALL *) | 49 (* TAIL_CALL *)
    | 144 (* STR_CONCAT *) ->
      if !ip < len then begin
        let v = bc.(!ip) in ip := !ip + 1;
        operands := [Number (float_of_int v)]
      end
    | 32 (* JUMP *) | 33 (* JUMP_IF_FALSE *) | 34 (* JUMP_IF_TRUE *) ->
      if !ip + 1 < len then begin
        let lo = bc.(!ip) in let hi = bc.(!ip + 1) in
        let raw = lo lor (hi lsl 8) in
        let signed = if raw >= 32768 then raw - 65536 else raw in
        ip := !ip + 2;
        operands := [Number (float_of_int signed)]
      end
    | 52 (* CALL_PRIM *) ->
      if !ip + 2 < len then begin
        let lo = bc.(!ip) in let hi = bc.(!ip + 1) in
        let idx = lo lor (hi lsl 8) in
        let argc = bc.(!ip + 2) in
        ip := !ip + 3;
        let prim_name = if idx < Array.length consts
          then (match consts.(idx) with String s -> s | _ -> "?") else "?" in
        operands := [Number (float_of_int idx); String prim_name; Number (float_of_int argc)]
      end
    | _ -> ());
    let entry = Hashtbl.create 4 in
    Hashtbl.replace entry "offset" (Number (float_of_int offset));
    Hashtbl.replace entry "opcode" (String name);
    Hashtbl.replace entry "operands" (List !operands);
    instrs := Dict entry :: !instrs
  done;
  let result = Hashtbl.create 4 in
  Hashtbl.replace result "arity" (Number (float_of_int code.vc_arity));
  Hashtbl.replace result "num_locals" (Number (float_of_int code.vc_locals));
  Hashtbl.replace result "constants" (List (Array.to_list (Array.map (fun v -> String (Sx_types.inspect v)) consts)));
  Hashtbl.replace result "bytecode" (List (List.rev !instrs));
  Dict result