(** 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 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 (Number n) -> int_of_float n | _ -> 0
    in
    let rest_arity = match find2 "rest-arity" "vc-rest-arity" with
      | 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;
    pop vm
  | 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
            | Number x, Number y -> Number (x +. 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
            | Number x, Number y -> Number (x -. 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
            | Number x, Number y -> Number (x *. 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
            | Number x, Number y -> Number (x /. y)
            | _ -> (Hashtbl.find Sx_primitives.primitives "/") [a; b])
        | 164 (* OP_EQ *) ->
          let b = pop vm and a = pop vm in
          let rec norm = function
            | ListRef { contents = l } -> List (List.map norm l)
            | List l -> List (List.map norm l) | v -> v in
          push vm (Bool (norm a = norm 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)
            | _ -> (Hashtbl.find Sx_primitives.primitives "<") [a; b])
        | 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)
            | _ -> (Hashtbl.find Sx_primitives.primitives ">") [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 } -> 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
            | _ -> (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;
      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