Merge pull request #1999 from nicolasnoble/fix-coroutine-lifetime

nicolasnoble · web-flow · commit d6e7b1e822a0 · 2026-04-08T14:55:47.000-07:00
Fix coroutine lifetime and realloc fragmentation
diff --git a/src/mips/openbios/kernel/alloc.c b/src/mips/openbios/kernel/alloc.c
@@ -238,7 +238,7 @@ static __attribute__((section(".ramtext"))) void *multi_realloc(void *ptr_, size
     if (block < head) {
         if (size < old_size) {
             empty_block *new_block = (empty_block *)((char *)block + size);
-            if (head == (empty_block *)((char *)block + size)) {
+            if (head == (empty_block *)((char *)block + old_size)) {
                 new_block->next = head->next;
                 new_block->size = head->size + (old_size - size);
             } else {
@@ -292,7 +292,7 @@ static __attribute__((section(".ramtext"))) void *multi_realloc(void *ptr_, size
 
         if (size < old_size) {
             empty_block *new_block = (empty_block *)((char *)block + size);
-            if ((next != &marker) && (((char *)block + size) == (char *)next)) {
+            if ((next != &marker) && (((char *)block + old_size) == (char *)next)) {
                 new_block->next = next->next;
                 new_block->size = old_size - size + next->size;
             } else {
diff --git a/src/mips/psyqo/coroutine.hh b/src/mips/psyqo/coroutine.hh
@@ -128,7 +128,7 @@ struct Coroutine {
      * @details This method is used to create an instance of the `Awaiter` object.
      * It's used to suspend the coroutine after scheduling an asynchronous operation.
      */
-    Awaiter awaiter() { return Awaiter(this); }
+    Awaiter awaiter() & { return Awaiter(this); }
 
     /**
      * @brief Resumes the coroutine.
@@ -232,18 +232,52 @@ struct Coroutine {
   public:
     using promise_type = Promise;
 
-    constexpr bool await_ready() { return m_handle.done(); }
-    template <typename U>
-    constexpr void await_suspend(std::coroutine_handle<U> h) {
-        m_handle.promise().m_awaitingCoroutine = h;
-        resume();
-    }
-    constexpr T await_resume() {
-        if constexpr (std::is_void<T>::value) {
-            return;
-        } else {
-            return eastl::move(m_handle.promise().m_value);
+    /**
+     * @brief The awaiter type for coroutine-to-coroutine chaining via co_await.
+     *
+     * @details GCC has known bugs across versions 10-15+ where the awaitable temporary
+     * in a `co_await` expression is not correctly promoted to the coroutine frame.
+     * By using `operator co_await()`, we transfer the coroutine handle to this
+     * awaiter object that GCC's coroutine pass reliably stores in the frame.
+     */
+    struct ChainAwaiter {
+        std::coroutine_handle<Promise> handle;
+
+        explicit ChainAwaiter(std::coroutine_handle<Promise> h) : handle(h) {}
+        ~ChainAwaiter() {
+            if (handle) handle.destroy();
+        }
+
+        ChainAwaiter(ChainAwaiter &&other) : handle(other.handle) { other.handle = nullptr; }
+        ChainAwaiter &operator=(ChainAwaiter &&) = delete;
+        ChainAwaiter(const ChainAwaiter &) = delete;
+        ChainAwaiter &operator=(const ChainAwaiter &) = delete;
+
+        constexpr bool await_ready() { return handle.done(); }
+
+        void await_suspend(std::coroutine_handle<> h) {
+            handle.promise().m_awaitingCoroutine = h;
+            if (!handle.done()) handle.resume();
+        }
+
+        constexpr T await_resume() {
+            if constexpr (std::is_void<T>::value) {
+                handle.destroy();
+                handle = nullptr;
+                return;
+            } else {
+                auto val = eastl::move(handle.promise().m_value);
+                handle.destroy();
+                handle = nullptr;
+                return val;
+            }
         }
+    };
+
+    ChainAwaiter operator co_await() && {
+        auto h = m_handle;
+        m_handle = nullptr;
+        return ChainAwaiter{h};
     }
 };
 
diff --git a/src/mips/psyqo/src/alloc.c b/src/mips/psyqo/src/alloc.c
@@ -584,8 +584,8 @@ void *psyqo_realloc(void *ptr_, size_t size_) {
         if (size < old_size) {
             // We're going to create a new block at the end of what we are re-allocating.
             empty_block *new_block = (empty_block *)((char *)block + size);
-            // Is the next block adjacent to the block we're re-allocating?
-            if ((next != &marker) && (((char *)block + size) == (char *)next)) {
+            // Is the next block adjacent to the end of our original allocation?
+            if ((next != &marker) && (((char *)block + old_size) == (char *)next)) {
                 // Yes, we can merge them.
                 new_block->next = next->next;
                 new_block->size = old_size - size + next->size;
diff --git a/src/support/coroutine.h b/src/support/coroutine.h
@@ -26,52 +26,86 @@ SOFTWARE.
 
 #pragma once
 
-#if defined(__APPLE__) && (__clang_major__ < 15)
-// Why has Apple become the Microsoft of Software Engineering?
-#include <experimental/coroutine>
-#else
 #include <coroutine>
-#endif
 #include <type_traits>
+#include <utility>
 
 namespace PCSX {
 
 template <typename T = void>
 struct Coroutine {
-#if defined(__APPLE__) && (__clang_major__ < 15)
-    template <typename U>
-    using CoroutineHandle = std::experimental::coroutine_handle<U>;
-    using CoroutineHandleVoid = std::experimental::coroutine_handle<void>;
-#else
-    template <typename U>
-    using CoroutineHandle = std::coroutine_handle<U>;
-    using CoroutineHandleVoid = std::coroutine_handle<void>;
-#endif
     struct Empty {};
     typedef typename std::conditional<std::is_void<T>::value, Empty, T>::type SafeT;
 
     Coroutine() = default;
-    Coroutine(Coroutine &&other) = default;
-    Coroutine &operator=(Coroutine &&other) = default;
+
+    Coroutine(Coroutine &&other)
+        : m_handle(std::exchange(other.m_handle, nullptr)),
+          m_value(std::move(other.m_value)),
+          m_suspended(other.m_suspended),
+          m_earlyResume(other.m_earlyResume) {
+        other.m_suspended = true;
+        other.m_earlyResume = false;
+    }
+
+    Coroutine &operator=(Coroutine &&other) {
+        if (this != &other) {
+            if (m_handle) m_handle.destroy();
+            m_handle = std::exchange(other.m_handle, nullptr);
+            m_value = std::move(other.m_value);
+            m_suspended = other.m_suspended;
+            m_earlyResume = other.m_earlyResume;
+            other.m_suspended = true;
+            other.m_earlyResume = false;
+        }
+        return *this;
+    }
+
     Coroutine(Coroutine const &) = delete;
     Coroutine &operator=(Coroutine const &) = delete;
+    ~Coroutine() {
+        if (m_handle) m_handle.destroy();
+        m_handle = nullptr;
+    }
 
     struct Awaiter {
-        constexpr bool await_ready() const noexcept { return false; }
-        constexpr void await_suspend(CoroutineHandleVoid) const noexcept {}
+        Awaiter(Awaiter &&other) = default;
+        Awaiter &operator=(Awaiter &&other) = default;
+        Awaiter(Awaiter const &) = default;
+        Awaiter &operator=(Awaiter const &) = default;
+        constexpr bool await_ready() const noexcept {
+            bool ret = m_coroutine->m_earlyResume;
+            m_coroutine->m_earlyResume = false;
+            return ret;
+        }
+        constexpr void await_suspend(std::coroutine_handle<> h) { m_coroutine->m_suspended = true; }
         constexpr void await_resume() const noexcept {}
+
+      private:
+        Awaiter(Coroutine *coroutine) : m_coroutine(coroutine) {}
+        Coroutine *m_coroutine;
+        friend struct Coroutine;
     };
 
-    Awaiter awaiter() { return {}; }
+    Awaiter awaiter() & { return Awaiter(this); }
+
     void resume() {
         if (!m_handle) return;
+        if (!m_suspended) {
+            m_earlyResume = true;
+            return;
+        }
+        m_suspended = false;
         m_handle.resume();
     }
 
     bool done() {
         if (!m_handle) return true;
         bool isDone = m_handle.done();
         if (isDone) {
+            if constexpr (!std::is_void<T>::value) {
+                m_value = std::move(m_handle.promise().m_value);
+            }
             m_handle.destroy();
             m_handle = nullptr;
         }
@@ -82,38 +116,106 @@ struct Coroutine {
 
   private:
     struct PromiseVoid {
-        Coroutine<T> get_return_object() { return Coroutine{std::move(CoroutineHandle<Promise>::from_promise(*this))}; }
-        Awaiter initial_suspend() { return {}; }
-        Awaiter final_suspend() noexcept { return {}; }
+        struct FinalAwaiter {
+            bool await_ready() const noexcept { return false; }
+            std::coroutine_handle<> await_suspend(std::coroutine_handle<PromiseVoid> h) const noexcept {
+                auto next = h.promise().m_awaitingCoroutine;
+                return next ? next : std::noop_coroutine();
+            }
+            void await_resume() const noexcept {}
+        };
+        Coroutine<> get_return_object() {
+            return Coroutine<>{std::move(std::coroutine_handle<Promise>::from_promise(*this))};
+        }
+        std::suspend_always initial_suspend() { return {}; }
+        FinalAwaiter final_suspend() noexcept { return {}; }
         void unhandled_exception() {}
         template <typename From>
-        Awaiter yield_value(From &&from) {
-            return {};
+        From yield_value(From &&from) {
+            return std::forward<From>(from);
         }
         void return_void() {}
+        [[no_unique_address]] Empty m_value;
+        std::coroutine_handle<> m_awaitingCoroutine;
     };
+
     struct PromiseValue {
-        PromiseValue(Coroutine<T> *c) : coroutine(c) {}
-        Coroutine<T> get_return_object() { return Coroutine{std::move(CoroutineHandle<Promise>::from_promise(*this))}; }
-        Awaiter initial_suspend() { return {}; }
-        Awaiter final_suspend() noexcept { return {}; }
+        struct FinalAwaiter {
+            bool await_ready() const noexcept { return false; }
+            std::coroutine_handle<> await_suspend(std::coroutine_handle<PromiseValue> h) const noexcept {
+                auto next = h.promise().m_awaitingCoroutine;
+                return next ? next : std::noop_coroutine();
+            }
+            void await_resume() const noexcept {}
+        };
+        Coroutine<T> get_return_object() {
+            return Coroutine{std::move(std::coroutine_handle<Promise>::from_promise(*this))};
+        }
+        std::suspend_always initial_suspend() { return {}; }
+        FinalAwaiter final_suspend() noexcept { return {}; }
         void unhandled_exception() {}
-        // This should be an std::convertible_to<T>, but Apple still doesn't have a fully C++-20 conformant library.
         template <typename From>
-        Awaiter yield_value(From &&from) {
-            coroutine->m_value = std::forward<From>(from);
-            return {};
+        From yield_value(From &&from) {
+            return std::forward<From>(from);
+        }
+        void return_value(T &&value) {
+            m_value = std::move(value);
         }
-        void return_value(T &&value) { coroutine->m_value = std::forward(value); }
-        Coroutine<T> *coroutine = nullptr;
+        T m_value;
+        std::coroutine_handle<> m_awaitingCoroutine;
     };
+
     typedef typename std::conditional<std::is_void<T>::value, PromiseVoid, PromiseValue>::type Promise;
-    Coroutine(CoroutineHandle<Promise> &&handle) : m_handle(std::move(handle)) {}
-    CoroutineHandle<Promise> m_handle;
+
+    Coroutine(std::coroutine_handle<Promise> &&handle) : m_handle(std::move(handle)) {}
+
+    std::coroutine_handle<Promise> m_handle;
     [[no_unique_address]] SafeT m_value;
+    bool m_suspended = true;
+    bool m_earlyResume = false;
 
   public:
     using promise_type = Promise;
+
+    struct ChainAwaiter {
+        std::coroutine_handle<Promise> handle;
+
+        explicit ChainAwaiter(std::coroutine_handle<Promise> h) : handle(h) {}
+        ~ChainAwaiter() {
+            if (handle) handle.destroy();
+        }
+
+        ChainAwaiter(ChainAwaiter &&other) : handle(other.handle) { other.handle = nullptr; }
+        ChainAwaiter &operator=(ChainAwaiter &&) = delete;
+        ChainAwaiter(const ChainAwaiter &) = delete;
+        ChainAwaiter &operator=(const ChainAwaiter &) = delete;
+
+        constexpr bool await_ready() { return handle.done(); }
+
+        void await_suspend(std::coroutine_handle<> h) {
+            handle.promise().m_awaitingCoroutine = h;
+            if (!handle.done()) handle.resume();
+        }
+
+        constexpr T await_resume() {
+            if constexpr (std::is_void<T>::value) {
+                handle.destroy();
+                handle = nullptr;
+                return;
+            } else {
+                auto val = std::move(handle.promise().m_value);
+                handle.destroy();
+                handle = nullptr;
+                return val;
+            }
+        }
+    };
+
+    ChainAwaiter operator co_await() && {
+        auto h = m_handle;
+        m_handle = nullptr;
+        return ChainAwaiter{h};
+    }
 };
 
 }  // namespace PCSX
