Merge branch 'loops/sx-vm-extensions' into loops/host

hosts/ocaml/lib/sx_vm.ml

@@ -336,30 +336,51 @@ and call_closure_reuse cl args =
     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
+      (try run vm;
+        (* Normal completion: result sits at the top of the stack.
+           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. *)
+        if vm.sp > saved_sp then vm.stack.(vm.sp - 1) else Nil
+       with
+       | VmSuspended (req, _) as e ->
+         (match !Sx_types._cek_io_resolver with
+          | Some resolver ->
+            (* Serving path: a `perform` fired inside this HO-primitive
+               callback (map/filter/reduce/for-each/…). The primitive's native
+               OCaml loop sits between us and the resume point, so we CANNOT
+               unwind it and resume later (the loop state would be lost and the
+               remaining elements dropped — corrupting the stack so the next
+               CALL_PRIM sees wrong args). Instead resolve the callback's IO
+               inline and run it to completion right here, returning its value
+               to the native loop exactly as a non-suspending callback would.
+               reuse_stack is isolated so an outer suspension's saved
+               continuations aren't consumed by this nested resume. *)
+            let saved_reuse = vm.reuse_stack in
+            vm.reuse_stack <- [];
+            let rec settle req =
+              let r = resolver req Nil in
+              (try resume_vm vm r
+               with VmSuspended (req2, _) -> settle req2)
+            in
+            let cb = settle req in
+            vm.reuse_stack <- saved_reuse;
+            cb
+          | None ->
+            (* CEK-driven path (no synchronous resolver): preserve the existing
+               behaviour — save the caller's continuation on the reuse stack and
+               re-raise so resume_vm restores it after the callback finishes.
+               DON'T merge frames — that corrupts the frame chain. *)
+            vm.reuse_stack <- (saved_frames, saved_sp) :: vm.reuse_stack;
+            raise e)
+       | e ->
+         vm.frames <- saved_frames;
+         vm.sp <- saved_sp;
+         raise e)
     in
+    vm.frames <- saved_frames;
     vm.sp <- saved_sp;
     result
   | None ->
@@ -938,7 +959,17 @@ and run vm =
 
     After the callback finishes, restores any call_closure_reuse
     continuations saved on vm.reuse_stack (innermost first). *)
-let resume_vm vm result =
+and resume_vm vm result =
+  (* The resumed execution runs on [vm]; HO primitives (map/filter/…) called
+     during the resume reach for [!_active_vm] to run their callbacks on the
+     same stack. call_closure restored [_active_vm] to the *caller* when the
+     original VmSuspended unwound through it, so without re-asserting it here
+     the resumed run's callbacks land on the wrong VM (or allocate a fresh
+     one), corrupting the stack. Mirror call_closure's save/set/restore. *)
+  let prev_active = !_active_vm in
+  _active_vm := Some vm;
+  let restore () = _active_vm := prev_active in
+  (try
   (match vm.pending_cek with
    | Some cek_state ->
      vm.pending_cek <- None;
@@ -1010,7 +1041,9 @@ let resume_vm vm result =
   let pending = List.rev vm.reuse_stack in
   vm.reuse_stack <- [];
   restore_reuse pending;
-  pop vm
+  let r = pop vm in
+  restore (); r
+  with e -> restore (); raise e)
 
 (** Execute a compiled module (top-level bytecode). *)
 let execute_module code globals =