fix(runtime): unwind all callers when guest code yields to the EE scheduler

ps2xRuntime/include/runtime/ee_scheduler.h

@@ -282,6 +282,12 @@ class EeScheduler
     void accountCycles(uint32_t cycles) noexcept;
     [[nodiscard]] bool isExecutingGuest() const noexcept;
 
+    // Guest code gives control back to the scheduler by returning up the host stack (a due checkpoint, or a
+    // branch the dispatcher cannot follow in place). Callers of a nested guest function compare this counter
+    // across the call to tell such an unwind apart from a normal return.
+    void noteGuestUnwind() noexcept { m_guestUnwindCount.fetch_add(1u, std::memory_order_relaxed); }
+    [[nodiscard]] uint64_t guestUnwindCount() const noexcept { return m_guestUnwindCount.load(std::memory_order_relaxed); }
+
     // Kernel object API. All calls except postEvent/requestStop execute on the
     // EE executor and therefore need no host synchronization.
     void setupCurrentThread(uint32_t stack, uint32_t stackSize, uint32_t gp);
@@ -430,6 +436,7 @@ class EeScheduler
     std::thread::id m_executorThread{};
     std::atomic<bool> m_running{false};
     std::atomic<bool> m_guestExecuting{false};
+    std::atomic<uint64_t> m_guestUnwindCount{0u};
     std::atomic<bool> m_stopRequested{false};
     std::atomic<bool> m_checkpointPending{false};
     uint32_t m_debugPublishCountdown = 0u;

ps2xRuntime/src/lib/ps2_runtime.cpp

@@ -1358,6 +1358,7 @@ bool PS2Runtime::dispatchGuestBranch(uint8_t *rdram,
     // this charge bounds straight-line call chains that have no local loop.
     if (m_eeScheduler && m_eeScheduler->checkpointDue(EeScheduler::kGuestDispatchCycles))
     {
+        m_eeScheduler->noteGuestUnwind();
         return false;
     }
 
@@ -1369,6 +1370,10 @@ bool PS2Runtime::dispatchGuestBranch(uint8_t *rdram,
         }
 
         ctx->pc = targetPc;
+        if (m_eeScheduler)
+        {
+            m_eeScheduler->noteGuestUnwind();
+        }
         return false;
     }
 
@@ -1396,19 +1401,38 @@ bool PS2Runtime::dispatchGuestBranch(uint8_t *rdram,
 
     RecompiledFunction targetFn = lookupFunction(targetPc);
     const uint32_t entryPc = ctx->pc;
+    const uint64_t unwindCountBefore = m_eeScheduler ? m_eeScheduler->guestUnwindCount() : 0u;
     targetFn(rdram, ctx, this);
 
     if (isStopRequested() || ctx->pc == 0u)
     {
         return false;
     }
 
+    // A callee that yielded to the scheduler can leave ctx->pc equal to its own entry (a recursive call, or a loop
+    // head on its first instruction), which is indistinguishable from an implicit return. The yield is recorded by
+    // the scheduler instead, and every caller on the host stack has to unwind too.
+    if (m_eeScheduler && m_eeScheduler->guestUnwindCount() != unwindCountBefore)
+    {
+        return false;
+    }
+
     if (ctx->pc == entryPc)
     {
         ctx->pc = fallthroughPc;
     }
 
-    return ctx->pc == fallthroughPc;
+    if (ctx->pc != fallthroughPc)
+    {
+        // Non-local return (longjmp, tail jump): the scheduler continues at ctx->pc.
+        if (m_eeScheduler)
+        {
+            m_eeScheduler->noteGuestUnwind();
+        }
+        return false;
+    }
+
+    return true;
 }
 
 void PS2Runtime::SignalException(R5900Context *ctx, PS2Exception exception)
@@ -2211,7 +2235,12 @@ void PS2Runtime::postEeEvent(EeEvent event)
 
 bool PS2Runtime::eeCheckpointDue(uint32_t cycles) noexcept
 {
-    return m_eeScheduler->checkpointDue(cycles);
+    const bool due = m_eeScheduler->checkpointDue(cycles);
+    if (due)
+    {
+        m_eeScheduler->noteGuestUnwind();
+    }
+    return due;
 }
 
 [[noreturn]] void PS2Runtime::eeWaitVSyncTicks(uint32_t ticks, uint32_t resumePc)

ps2xTest/src/ps2_runtime_expansion_tests.cpp

@@ -166,6 +166,53 @@ namespace
         }
     }
 
+    // Makes the next scheduler checkpoint due, like a timer or vblank event arriving during guest code.
+    void requestSchedulerCheckpoint(PS2Runtime *runtime)
+    {
+        runtime->postEeEvent(EeEvent{EeEventType::ExternalWake, 0u, 0u});
+    }
+
+    // Generated code for a recursive call: the callee (same entry as this function) is dispatched, finds a
+    // checkpoint due, and leaves ctx->pc at its own entry; the caller then returns up the host stack.
+    void testGuestBranchRecursiveYieldHandler(uint8_t *rdram, R5900Context *ctx, PS2Runtime *runtime)
+    {
+        if (!ctx || !runtime)
+        {
+            return;
+        }
+
+        requestSchedulerCheckpoint(runtime);
+        if (!runtime->dispatchGuestBranch(rdram,
+                                          ctx,
+                                          0x3300u,
+                                          0x3304u,
+                                          0x3308u,
+                                          PS2Runtime::GuestBranchKind::DirectCall,
+                                          "test-recursive-call"))
+        {
+            return;
+        }
+        setRegU32(*ctx, 2, 0x00BAD001u);
+    }
+
+    // A function whose first instruction is a loop head: the backward branch finds a checkpoint due, so the
+    // function returns with ctx->pc still equal to its entry.
+    void testGuestBranchEntryLoopYieldHandler(uint8_t *, R5900Context *ctx, PS2Runtime *runtime)
+    {
+        if (!ctx || !runtime)
+        {
+            return;
+        }
+
+        requestSchedulerCheckpoint(runtime);
+        ctx->pc = 0x3400u;
+        if (runtime->eeCheckpointDue())
+        {
+            return;
+        }
+        setRegU32(*ctx, 2, 0x00BAD002u);
+    }
+
     std::atomic<uint32_t> gGuestJumpTargetCount{0u};
 
     void testGuestJumpTargetHandler(uint8_t *, R5900Context *, PS2Runtime *)
@@ -455,6 +502,56 @@ void register_ps2_runtime_expansion_tests()
                      "callee transfer PC should be preserved");
         });
 
+        tc.Run("dispatchGuestBranch does not treat a recursive yield as a return", [](TestCase &t)
+        {
+            PS2Runtime runtime;
+            runtime.registerFunction(0x3300u, &testGuestBranchRecursiveYieldHandler);
+
+            R5900Context ctx{};
+            ctx.pc = 0x2000u;
+
+            const bool returnedToFallthrough = runtime.dispatchGuestBranch(
+                nullptr,
+                &ctx,
+                0x3300u,
+                0x2000u,
+                0x2008u,
+                PS2Runtime::GuestBranchKind::DirectCall,
+                "test-yield-outer");
+
+            t.IsFalse(returnedToFallthrough,
+                      "a yield inside a nested call must unwind every caller instead of returning to its fallthrough");
+            t.Equals(ctx.pc, 0x3300u,
+                     "the yielded callee entry should stay in ctx->pc for the scheduler to resume");
+            t.Equals(::getRegU32(&ctx, 2), 0u,
+                     "no caller should continue running after the yield");
+        });
+
+        tc.Run("dispatchGuestBranch does not treat a yield at the callee entry as a return", [](TestCase &t)
+        {
+            PS2Runtime runtime;
+            runtime.registerFunction(0x3400u, &testGuestBranchEntryLoopYieldHandler);
+
+            R5900Context ctx{};
+            ctx.pc = 0x2000u;
+
+            const bool returnedToFallthrough = runtime.dispatchGuestBranch(
+                nullptr,
+                &ctx,
+                0x3400u,
+                0x2000u,
+                0x2008u,
+                PS2Runtime::GuestBranchKind::DirectCall,
+                "test-entry-loop-yield");
+
+            t.IsFalse(returnedToFallthrough,
+                      "a checkpoint yield that leaves ctx->pc at the callee entry must not look like a return");
+            t.Equals(ctx.pc, 0x3400u,
+                     "the callee entry should stay in ctx->pc for the scheduler to resume");
+            t.Equals(::getRegU32(&ctx, 2), 0u,
+                     "the callee should not continue past the yield");
+        });
+
         tc.Run("dispatchGuestBranch rejects missing exact targets", [](TestCase &t)
         {
             PS2Runtime runtime;