add event, fix shared-task

Author: denisichh

runtime-light/coroutine/async-stack-methods.h

@@ -10,6 +10,7 @@
 #include <utility>
 
 #include "runtime-light/coroutine/coroutine-state.h"
+#include "runtime-light/stdlib/diagnostics/logs.h"
 
 namespace kphp::coro {
 
@@ -19,11 +20,13 @@ namespace kphp::coro {
  * capturing one of the stack frames in the synchronous stack trace and jump between parts of the async_stack.
  */
 [[clang::noinline]] inline void resume_with_new_root(std::coroutine_handle<> handle, async_stack_root* new_async_stack_root) noexcept {
+  kphp::log::assertion(new_async_stack_root != nullptr);
   auto& coroutine_st{CoroutineInstanceState::get()};
-  auto* previous_stack_root = coroutine_st.current_async_stack_root;
-  new_async_stack_root->next_async_stack_root = coroutine_st.current_async_stack_root;
+  auto* previous_stack_root{coroutine_st.current_async_stack_root};
 
+  new_async_stack_root->next_async_stack_root = coroutine_st.current_async_stack_root;
   new_async_stack_root->stop_sync_stack_frame = reinterpret_cast<stack_frame*>(STACK_FRAME_ADDRESS);
+
   coroutine_st.current_async_stack_root = new_async_stack_root;
   handle.resume();
   coroutine_st.current_async_stack_root = previous_stack_root;

runtime-light/coroutine/detail/when-any.h

@@ -194,7 +194,6 @@ class when_any_task_promise_base : public kphp::coro::async_stack_element {
   auto start(when_any_latch& latch, void* return_address) noexcept {
     m_latch = std::addressof(latch);
 
-
     kphp::coro::async_stack_root root{};
     auto& async_stack_frame{get_async_stack_frame()};
     // initialize when_all_task's async stack frame and make it the top frame

runtime-light/coroutine/event.h

@@ -24,15 +24,15 @@ class event {
     event& m_event;
     bool m_suspended{};
     std::coroutine_handle<> m_awaiting_coroutine;
-    kphp::coro::async_stack_root& m_async_stack_root;
+    kphp::coro::async_stack_root* m_async_stack_root;
     kphp::coro::async_stack_frame* m_caller_async_stack_frame{};
 
     awaiter* m_next{};
     awaiter* m_prev{};
 
     explicit awaiter(event& event) noexcept
         : m_event(event),
-          m_async_stack_root(CoroutineInstanceState::get().base_coroutine_stack_root) {}
+          m_async_stack_root(CoroutineInstanceState::get().current_async_stack_root) {}
 
     awaiter(const awaiter&) = delete;
     awaiter(awaiter&&) = delete;
@@ -85,7 +85,7 @@ inline auto event::awaiter::await_ready() const noexcept -> bool {
 template<std::derived_from<kphp::coro::async_stack_element> caller_promise_type>
 auto event::awaiter::await_suspend(std::coroutine_handle<caller_promise_type> awaiting_coroutine) noexcept -> void {
   // save caller's async stack frame
-  m_caller_async_stack_frame = m_async_stack_root.top_async_stack_frame;
+  m_caller_async_stack_frame = m_async_stack_root->top_async_stack_frame;
 
   m_suspended = true;
   m_awaiting_coroutine = awaiting_coroutine;
@@ -105,7 +105,7 @@ inline auto event::awaiter::await_resume() noexcept -> void {
   if (std::exchange(m_suspended, false)) {
     // restore caller's async stack frame if it was suspended
     kphp::log::assertion(m_caller_async_stack_frame != nullptr);
-    m_async_stack_root.top_async_stack_frame = std::exchange(m_caller_async_stack_frame, nullptr);
+    m_async_stack_root->top_async_stack_frame = std::exchange(m_caller_async_stack_frame, nullptr);
   }
 }
 

runtime-light/coroutine/shared-task.h

@@ -52,8 +52,12 @@ struct promise_base : kphp::coro::async_stack_element {
           // read the m_next pointer before resuming the coroutine
           // since resuming the coroutine may destroy the shared_task_waiter value
           auto* next{awaiter->m_next};
-          auto* async_stack_root{promise.get_async_stack_frame().async_stack_root};
-          kphp::coro::resume_with_new_root(awaiter->m_continuation, async_stack_root);
+          auto& async_stack_frame{promise.get_async_stack_frame()};
+
+          kphp::coro::async_stack_root root{};
+          root.top_async_stack_frame = std::addressof(async_stack_frame);
+          async_stack_frame.async_stack_root = std::addressof(root);
+          kphp::coro::resume_with_new_root(awaiter->m_continuation, std::addressof(root));
           awaiter = next;
         }
         // return last awaiter's coroutine_handle to allow it to potentially be compiled as a tail-call
@@ -100,8 +104,12 @@ struct promise_base : kphp::coro::async_stack_element {
     if (m_awaiters == NOT_STARTED_VAL) {
       m_awaiters = STARTED_NO_WAITERS_VAL;
       const auto& handle{std::coroutine_handle<promise_type>::from_promise(*static_cast<promise_type*>(this))};
-      auto* async_stack_root{get_async_stack_frame().async_stack_root};
-      kphp::coro::resume_with_new_root(handle, async_stack_root);
+      auto& async_stack_frame{get_async_stack_frame()};
+
+      kphp::coro::async_stack_root root{};
+      root.top_async_stack_frame = std::addressof(async_stack_frame);
+      async_stack_frame.async_stack_root = std::addressof(root);
+      kphp::coro::resume_with_new_root(handle, std::addressof(root));
     }
     // coroutine already completed, don't suspend
     if (done()) {
@@ -169,24 +177,13 @@ template<typename promise_type>
 class awaiter_base {
   bool m_suspended{};
 
-  void set_async_top_frame(async_stack_frame& caller_frame, void* return_address) noexcept {
+  void set_async_top_frame(void* return_address) noexcept {
     /**
      * shared_task is the top of the stack for calls from it.
      * Therefore, it's awaiter doesn't store caller_frame, but it save `await_suspend()` return address
      * */
     async_stack_frame& callee_frame{m_coro.promise().get_async_stack_frame()};
-
     callee_frame.return_address = return_address;
-    auto* async_stack_root{caller_frame.async_stack_root};
-    kphp::log::assertion(async_stack_root != nullptr);
-    callee_frame.async_stack_root = async_stack_root;
-    async_stack_root->top_async_stack_frame = std::addressof(callee_frame);
-  }
-
-  void reset_async_top_frame(async_stack_frame& caller_frame) noexcept {
-    auto* async_stack_root{caller_frame.async_stack_root};
-    kphp::log::assertion(async_stack_root != nullptr);
-    async_stack_root->top_async_stack_frame = std::addressof(caller_frame);
   }
 
 protected:
@@ -218,10 +215,9 @@ class awaiter_base {
 
   template<std::derived_from<kphp::coro::async_stack_element> caller_promise_type>
   [[clang::noinline]] auto await_suspend(std::coroutine_handle<caller_promise_type> awaiting_coroutine) noexcept -> bool {
-    set_async_top_frame(awaiting_coroutine.promise().get_async_stack_frame(), STACK_RETURN_ADDRESS);
+    set_async_top_frame(STACK_RETURN_ADDRESS);
     m_awaiter.m_continuation = awaiting_coroutine;
     m_suspended = m_coro.promise().suspend_awaiter(m_awaiter);
-    reset_async_top_frame(awaiting_coroutine.promise().get_async_stack_frame());
     return m_suspended;
   }