From a4baff8ff5fc8ccfbf36a2ef0c4458419061b318 Mon Sep 17 00:00:00 2001
From: Khalil Estell <kammce0@gmail.com>
Date: Mon, 30 Mar 2026 16:57:05 -0700
Subject: [PATCH] :bug: Add support for operation stacking
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Operation stacking is when you load a context with multiple coroutines.
Before, this resulted in a memory leak, due to the active_handle being
overwritten and lost. This also results in lifetime violation if the
context stack, prior to the new coroutine being loaded, had objects
with lifetimes that need to be destroyed. This change enables operation
stacking so it is no longer erroneous behavior. Now the last coroutine
loaded is the first one to be executed (LIFO) and up the chain until
the first loaded coroutine is reached.

\## Summary

- Enables multiple coroutines to be loaded onto the same `context` in
  LIFO order — the last routine pushed runs first, allowing natural
  "stacking" of async work on a single context
- Bans co-awaiting an l-value `future` (deleted `operator co_await() &`)
  to prevent accidentally awaiting a future from a different context;
  co-awaiting a future whose context was allocated inside a coroutine
  frame is now a contract violation (`contract_assert`) or
  `std::terminate`
- Removes `[[clang::lifetimebound]]` from `awaiter` constructor
- Simplify `await_suspend` to just return the stored handle and drop the
  explicit continuation chain. Continuation is now captured at
  `get_return_object()` time via `m_context->active_handle()` to enable
  operation stacking.

\## Test plan

- [x] `tests/async_stacking.test.cpp` — LIFO ordering with two and three routines on the same context
- [x] `tests/context_swapping.test.cpp` — verifies that co-awaiting a future from a mismatched nested context triggers `std::terminate` (or contract violation when contracts are enabled)
- [x] Existing `tests/basics.test.cpp` continues to pass with the updated `await_ready` logic
---
 CMakeLists.txt                   |   3 +-
 modules/coroutine.cppm           |  91 +++++++++++++---
 tests/async_stacking.test.cpp    | 173 +++++++++++++++++++++++++++++++
 tests/basics.test.cpp            |   8 +-
 tests/basics_dep_inject.test.cpp | 134 ------------------------
 tests/context_swapping.test.cpp  |  88 ++++++++++++++++
 6 files changed, 342 insertions(+), 155 deletions(-)
 create mode 100644 tests/async_stacking.test.cpp
 delete mode 100644 tests/basics_dep_inject.test.cpp
 create mode 100644 tests/context_swapping.test.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6701ca0..e4f01bb 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -36,11 +36,12 @@ libhal_add_tests(async_context
         cancel
         exclusive_access
         proxy
-        basics_dep_inject
         on_unblock
         simple_scheduler
         clock_adapter
         run_until_done
+        async_stacking
+        context_swapping
 
     MODULES
         tests/util.cppm
diff --git a/modules/coroutine.cppm b/modules/coroutine.cppm
index 2a52edf..d576e72 100644
--- a/modules/coroutine.cppm
+++ b/modules/coroutine.cppm
@@ -221,8 +221,13 @@ export struct bad_coroutine_alloc : std::bad_alloc
  * has been cancelled. It indicates that the operation was explicitly cancelled
  * before completion.
  */
-export class operation_cancelled : public std::exception
+export struct operation_cancelled : public std::exception
 {
+  operation_cancelled(void const* p_future_address)
+    : future_address(p_future_address)
+  {
+  }
+
   /**
    * @brief Get exception message
    *
@@ -232,6 +237,8 @@ export class operation_cancelled : public std::exception
   {
     return "This future has been cancelled!";
   }
+
+  void const* future_address = nullptr;
 };
 
 // =============================================================================
@@ -771,7 +778,8 @@ public:
    * destroyed before it is removed from this context.
    *
    * @param p_listener - the address of the unblock listener to be invoked when
-   * this context is unblocked.
+   * this context is unblocked. A nullptr may be passed to this parameter. It
+   * has the same effect as calling `clear_unblock_listener()`.
    */
   void on_unblock(unblock_listener* p_listener)
   {
@@ -1926,6 +1934,11 @@ public:
     }
   }
 
+  constexpr bool cancelled()
+  {
+    return std::holds_alternative<cancelled_state>(m_state);
+  }
+
   constexpr void resume() const
   {
     if (std::holds_alternative<handle_type>(m_state)) {
@@ -2004,30 +2017,34 @@ public:
   {
     future<T>& m_operation;
 
-    constexpr explicit awaiter(future<T>& p_operation
-                               [[clang::lifetimebound]]) noexcept
+    constexpr explicit awaiter(future<T>& p_operation) noexcept
       : m_operation(p_operation)
     {
     }
 
     [[nodiscard]] constexpr bool await_ready() const noexcept
     {
-      return m_operation.m_state.index() >= 1;
+      return not std::holds_alternative<handle_type>(m_operation.m_state);
     }
 
+    /**
+     * @brief Communicates to the awaiter to simply resume this coroutine
+     * associated with this future.
+     *
+     * @tparam U - return type of the promise
+     * @param p_calling_coroutine - this type is forgotten. The parent calling
+     * coroutine's handle was captured when the future object was created via
+     * get_return_object().
+     * @return std::coroutine_handle<> - self to continue
+     */
     template<typename U>
     std::coroutine_handle<> await_suspend(
-      std::coroutine_handle<promise<U>> p_calling_coroutine) noexcept
+      [[maybe_unused]] std::coroutine_handle<promise<U>>
+        p_calling_coroutine) noexcept
     {
       // This will not throw because the discriminate check was performed in
-      // `await_ready()` via the done() function. `done()` checks if the state
-      // is `handle_type` and if it is, it returns false, causing the code to
-      // call await_suspend().
-      auto handle = std::get<handle_type>(m_operation.m_state);
-      std::coroutine_handle<promise<U>>::from_address(handle.address())
-        .promise()
-        .continuation(p_calling_coroutine);
-      return handle;
+      // `await_ready()`.
+      return std::get<handle_type>(m_operation.m_state);
     }
 
     [[nodiscard]] constexpr monostate_or<T>& await_resume() const
@@ -2039,9 +2056,21 @@ public:
                    m_operation.m_state)) [[unlikely]] {
         std::rethrow_exception(
           std::get<std::exception_ptr>(m_operation.m_state));
+      } else if (std::holds_alternative<cancelled_state>(m_operation.m_state))
+        [[unlikely]] {
+        throw operation_cancelled{ &m_operation };
       }
 
-      throw operation_cancelled{};
+      // In the event that this coroutine awaiting this awaitable has
+      // requested the result of this awaitable before it has finished, then:
+      //
+      // - If contracts are enabled, then contract violation handler is called.
+      // - Otherwise, std::terminate is called.
+#if defined(__cpp_contracts)
+      contract_assert(std::holds_alternative<handle_type>(m_operation.m_state));
+#else
+      std::terminate();
+#endif
     }
 
     constexpr void await_resume() const
@@ -2054,9 +2083,21 @@ public:
                    m_operation.m_state)) [[unlikely]] {
         std::rethrow_exception(
           std::get<std::exception_ptr>(m_operation.m_state));
+      } else if (std::holds_alternative<cancelled_state>(m_operation.m_state))
+        [[unlikely]] {
+        throw operation_cancelled{ &m_operation };
       }
 
-      throw operation_cancelled{};
+      // In the event that this coroutine awaiting this awaitable has
+      // requested the result of this awaitable before it has finished, then:
+      //
+      // - If contracts are enabled, then contract violation handler is called.
+      // - Otherwise, std::terminate is called.
+#if defined(__cpp_contracts)
+      contract_assert(std::holds_alternative<handle_type>(m_operation.m_state));
+#else
+      std::terminate();
+#endif
     }
   };
 
@@ -2069,15 +2110,27 @@ public:
    * @return awaiter - An awaiter object that handles the suspension and
    * resumption of coroutines awaiting this future.
    *
+   * @pre The coroutine's context and the future's context are the same.
    * @note The awaiter will suspend the calling coroutine until this future
    * completes, then resume with either the result value or an exception. The
    * future will never be cancelled.
    */
-  [[nodiscard]] constexpr awaiter operator co_await() noexcept
+  [[nodiscard]] constexpr awaiter operator co_await() && noexcept
   {
     return awaiter{ *this };
   }
 
+  /**
+   * @brief co_awaiting an l-value future is banned
+   *
+   * Co-awaiting an l-value future provides a means to accidentally await on a
+   * future with a different context than the awaiting coroutine. This is unsafe
+   * and thus, banned.
+   *
+   * @return Nothing, deleted implementation
+   */
+  [[nodiscard]] constexpr awaiter operator co_await() & noexcept = delete;
+
 private:
   friend promise_type;
 
@@ -2127,6 +2180,10 @@ constexpr future<T> promise<T>::get_return_object() noexcept
 {
   using future_handle = std::coroutine_handle<promise<T>>;
   auto handle = future_handle::from_promise(*this);
+  // Chain: whatever was active becomes our continuation.
+  // If nothing was active (noop_sentinel), this is a normal top-level
+  // coroutine. If something WAS active, we implicitly sit on top of it.
+  m_continuation = m_context->active_handle();
   m_context->active_handle(handle);
   return future<T>{ handle };
 }
diff --git a/tests/async_stacking.test.cpp b/tests/async_stacking.test.cpp
new file mode 100644
index 0000000..914de03
--- /dev/null
+++ b/tests/async_stacking.test.cpp
@@ -0,0 +1,173 @@
+#include <coroutine>
+
+#include <boost/ut.hpp>
+
+import async_context;
+import test_utils;
+
+// doesn't matter...
+// bar runs until completion then foo runs until complete
+void async_stacking()
+{
+  using namespace boost::ut;
+
+  // LIFO stacking: the last routine pushed onto the context is the first
+  // to execute. When two routines are loaded onto the same context, the
+  // second one loaded runs first and the first one loaded runs last.
+
+  "two routines lifo order two steps each"_test = []() {
+    // Setup
+    async::inplace_context<1024> ctx;
+
+    unsigned step = 0;
+
+    auto routine_a = [&step](async::context&) -> async::future<void> {
+      // routine_a is loaded first; it runs last (LIFO)
+      step = 3;
+      co_await std::suspend_always{};
+      step = 4;
+      co_await std::suspend_always{};
+      co_return;
+    };
+
+    auto routine_b = [&step](async::context&) -> async::future<void> {
+      // routine_b is loaded second; it runs first (LIFO)
+      step = 1;
+      co_await std::suspend_always{};
+      step = 2;
+      co_await std::suspend_always{};
+      co_return;
+    };
+
+    // Exercise: load routine_a first, then routine_b
+    auto future_a = routine_a(ctx);
+    auto future_b = routine_b(ctx);
+
+    // Verify: neither has started yet
+    expect(that % not future_a.done());
+    expect(that % not future_b.done());
+    expect(that % 0 == step);
+
+    // Exercise: first resume — routine_b (last loaded) runs first, hits step=1
+    // then suspends
+    future_b.resume();
+
+    // Verify: routine_b ran first and suspended at step 1
+    expect(that % not future_b.done());
+    expect(that % not future_a.done());
+    expect(that % 1 == step);
+
+    // Exercise: second resume — routine_b resumes, hits step=2, suspends again
+    future_b.resume();
+
+    // Verify: routine_b suspended at step 2; routine_a still waiting
+    expect(that % not future_b.done());
+    expect(that % not future_a.done());
+    expect(that % 2 == step);
+
+    // Exercise: third resume — routine_b hits co_return and completes
+    future_b.resume();
+
+    // Verify: routine_b is done; routine_a still waiting
+    expect(that % future_b.done());
+    expect(that % not future_a.done());
+    expect(that % 3 == step);
+
+    // Exercise: fourth resume — routine_a (first loaded) now runs, hits step=3,
+    // then suspends
+    future_a.resume();
+
+    // Verify: routine_a ran and suspended at step 3
+    expect(that % not future_a.done());
+    expect(that % 4 == step);
+
+    // Exercise: fifth resume — routine_a resumes, hits step=4, suspends again
+    future_a.resume();
+
+    // Verify: routine_a is done; all memory released
+    expect(that % future_a.done());
+    expect(that % 0 == ctx.memory_used());
+    expect(that % 4 == step);
+  };
+
+  "three routines lifo order"_test = []() {
+    // Setup
+    async::inplace_context<2048> ctx;
+
+    unsigned step = 0;
+
+    auto routine_a = [&step](async::context&) -> async::future<void> {
+      // loaded first — runs last
+      step = 7;
+      co_await std::suspend_always{};
+      step = 8;
+      co_await std::suspend_always{};
+      co_return;
+    };
+
+    auto routine_b = [&step](async::context&) -> async::future<void> {
+      // loaded second — runs second
+      step = 4;
+      co_await std::suspend_always{};
+      step = 5;
+      co_await std::suspend_always{};
+      co_return;
+    };
+
+    auto routine_c = [&step](async::context&) -> async::future<void> {
+      // loaded third — runs first (LIFO)
+      step = 1;
+      co_await std::suspend_always{};
+      step = 2;
+      co_await std::suspend_always{};
+      co_return;
+    };
+
+    // Load in order: a, b, c
+    auto future_a = routine_a(ctx);
+    auto future_b = routine_b(ctx);
+    auto future_c = routine_c(ctx);
+
+    expect(that % 0 == step);
+
+    // routine_c runs first
+    future_c.resume();
+    expect(that % 1 == step);
+    expect(that % not future_c.done());
+
+    future_c.resume();
+    expect(that % 2 == step);
+    expect(that % not future_c.done());
+
+    future_c.resume();
+    expect(that % 4 == step);
+    expect(that % future_c.done());
+
+    // routine_b was already started by routine_c's co_return; resume to step=5
+    future_b.resume();
+    expect(that % 5 == step);
+    expect(that % not future_b.done());
+
+    // routine_b completes and immediately starts routine_a which runs to step=7
+    future_b.resume();
+    expect(that % 7 == step);
+    expect(that % future_b.done());
+    expect(that % not future_a.done());
+
+    future_a.resume();
+    expect(that % 8 == step);
+    expect(that % not future_a.done());
+
+    future_a.resume();
+    expect(that % 8 == step);
+    expect(that % future_a.done());
+
+    // All memory should be released
+    expect(that % 0 == ctx.memory_used());
+  };
+}
+
+int main()
+{
+  async_stacking();
+}
diff --git a/tests/basics.test.cpp b/tests/basics.test.cpp
index c555673..37e9167 100644
--- a/tests/basics.test.cpp
+++ b/tests/basics.test.cpp
@@ -1,6 +1,8 @@
 #include <coroutine>
 
 #include <boost/ut.hpp>
+#include <print>
+#include <variant>
 
 import async_context;
 import test_utils;
@@ -164,9 +166,9 @@ void basics()
     };
     auto co = [&step, &co2](async::context& p_ctx) -> async::future<int> {
       step = 1;  // skipped as the co2 will immediately start
-      [[maybe_unused]] auto val = co_await co2(p_ctx);
+      auto const val = co_await co2(p_ctx);
       step = 4;
-      co_return expected_return_value;
+      co_return val;
     };
 
     // Exercise 1
@@ -198,7 +200,7 @@ void basics()
     expect(that % 4 == step);
   };
 
-  "co_await coroutine"_test = []() {
+  "co_await coroutine sync"_test = []() {
     // Setup
     async::inplace_context<1024> ctx;
 
diff --git a/tests/basics_dep_inject.test.cpp b/tests/basics_dep_inject.test.cpp
deleted file mode 100644
index 36061f1..0000000
--- a/tests/basics_dep_inject.test.cpp
+++ /dev/null
@@ -1,134 +0,0 @@
-#include <coroutine>
-#include <optional>
-
-#include <boost/ut.hpp>
-#include <utility>
-
-import async_context;
-
-void basics_dep_inject()
-{
-  using namespace boost::ut;
-
-  "sync return type void"_test = []() {
-    // Setup
-    std::array<async::stack_word, 1024> stack{};
-    async::context ctx{ stack };
-
-    unsigned step = 0;
-    auto sync_coroutine = [&step](async::context&) -> async::future<void> {
-      step = 1;
-      return {};
-    };
-
-    // Exercise
-    auto future = sync_coroutine(ctx);
-
-    // Verify
-    expect(that % 0 == ctx.memory_used());
-    expect(that % future.done());
-    expect(that % future.has_value());
-    expect(that % 1 == step);
-  };
-
-  "suspend then co_return"_test = []() {
-    // Setup
-    std::array<async::stack_word, 1024> stack{};
-    async::context ctx{ stack };
-
-    static constexpr int expected_return_value = 1413;
-    unsigned step = 0;
-    auto async_coroutine =
-      [&step](async::context& p_ctx) -> async::future<int> {
-      step = 1;
-      while (step != 2) {
-        // external set to 2
-        co_await p_ctx.block_by_io();
-      }
-      step = 2;
-      co_return expected_return_value;
-    };
-
-    // Exercise 1
-    auto future = async_coroutine(ctx);
-
-    // Verify 1
-    expect(that % 0 < ctx.memory_used());
-    expect(that % not future.done());
-    expect(that % not future.has_value());
-    expect(that % 0 == step);
-
-    // Exercise 2: start and suspend coroutine
-    ctx.resume();
-
-    // Verify 2
-    expect(that % 0 < ctx.memory_used());
-    expect(that % not future.done());
-    expect(that % not future.has_value());
-    expect(that % 1 == step);
-    expect(ctx.state() == async::blocked_by::io);
-
-    // Exercise 3: resume and co_return from coroutine
-    ctx.unblock();
-    step = 2;
-    ctx.resume();
-
-    // Verify 3
-    expect(that % 0 == ctx.memory_used());
-    expect(that % future.done());
-    expect(that % future.has_value());
-    expect(that % expected_return_value == future.value());
-    expect(that % 2 == step);
-  };
-
-  "Call handler"_test = []() {
-    // Setup
-    std::array<async::stack_word, 1024> stack{};
-    async::context ctx{ stack };
-
-    static constexpr int expected_return_value = 1413;
-    unsigned step = 0;
-    auto async_coroutine =
-      [&step](async::context& p_ctx) -> async::future<int> {
-      step = 1;
-      co_await p_ctx.block_by_io();
-      step = 2;
-      co_return expected_return_value;
-    };
-
-    // Exercise 1
-    auto future = async_coroutine(ctx);
-
-    // Verify 1
-    expect(that % 0 < ctx.memory_used());
-    expect(that % not future.done());
-    expect(that % not future.has_value());
-    expect(that % 0 == step);
-
-    // Exercise 2: start and suspend coroutine
-    ctx.resume();
-
-    // Verify 2
-    expect(that % 0 < ctx.memory_used());
-    expect(async::blocked_by::io == ctx.state());
-    // expect(that % async::block_by::io == *info);
-    expect(that % not future.done());
-    expect(that % not future.has_value());
-    expect(that % 1 == step);
-
-    // Exercise 3: resume and co_return from coroutine
-    ctx.resume();
-
-    // Verify 3
-    expect(that % 0 == ctx.memory_used());
-    expect(that % future.done());
-    expect(that % future.has_value());
-    expect(that % expected_return_value == future.value());
-    expect(that % 2 == step);
-  };
-};
-
-int main()
-{
-  basics_dep_inject();
-}
diff --git a/tests/context_swapping.test.cpp b/tests/context_swapping.test.cpp
new file mode 100644
index 0000000..9c219b4
--- /dev/null
+++ b/tests/context_swapping.test.cpp
@@ -0,0 +1,88 @@
+#include <coroutine>
+
+#include <boost/ut.hpp>
+
+import async_context;
+import test_utils;
+
+void context_swapping()
+{
+  using namespace boost::ut;
+
+  // Context swapping: one routine runs on an outer context and produces a
+  // future. A second routine on a separate context co_awaits that future using
+  // a locally nested context. Awaiting a future whose context was allocated
+  // inside a coroutine frame is a contract violation and will abort.
+
+  "co_await future from another context"_test = []() {
+    // Setup
+    async::inplace_context<1024> ctx;
+
+    unsigned step = 0;
+
+    // routine_a produces a future that routine will await
+    auto routine_a = [&step](async::context& p_ctx) -> async::future<int> {
+      step = 1;
+      co_await std::suspend_always{};
+      step = 2;
+      co_await std::suspend_always{};
+      step = 3;
+      co_return 42;
+    };
+
+    // routine captures routine_a and co_awaits it using a nested context
+    // allocated inside the coroutine frame — this is a contract violation
+    auto routine = [&step, &routine_a](async::context&) -> async::future<int> {
+      step = 10;
+      co_await std::suspend_always{};
+
+      async::inplace_context<128> nested_ctx;
+      // This is a bug and a contract violation, this will either call the
+      // contract violator OR call std::terminate
+      auto result = co_await routine_a(nested_ctx);
+      // These should never be reached
+      step = 11;
+      co_return result;
+    };
+
+    // Exercise: start routine on the outer context
+    auto future = routine(ctx);
+
+    // Verify: routine has not run yet
+    expect(that % not future.done());
+    expect(that % 0 == step);
+    // ctx holds routine's frame
+    expect(that % 0 < ctx.memory_used());
+
+    // Start routine — it runs to step=10 then suspends at the first
+    // co_await suspend_always
+    future.resume();
+    expect(that % 10 == step);
+    expect(that % not future.done());
+
+    // Resume routine — it allocates nested_ctx and starts routine_a,
+    // which runs to step=1 then suspends
+    future.resume();
+    expect(that % 1 == step);
+
+// Resume routine — it hits co_await on a future whose context was
+// allocated inside the coroutine frame; this is a contract violation
+#if defined(__cpp_contracts)
+    // TODO(kammce): Add violation handler
+#else
+    // NOTE: aborts() forks a child process, so a std::set_terminate handler
+    // set here cannot be observed from the parent — terminate_was_called would
+    // always be false regardless. The aborts() check is the only observable
+    // assertion available across the process boundary.
+    expect(aborts([&future]() { future.resume(); }));
+    expect(that % 1 == step);
+    expect(that % not future.done());
+    // ctx holds routine's suspended frame while the nested context does work
+#endif
+  };
+}
+
+int main()
+{
+  context_swapping();
+}