From 4db885607a3927be7517d42fe2c6e30128c957d7 Mon Sep 17 00:00:00 2001
From: Ian Petersen <ispeters@gmail.com>
Date: Wed, 28 Aug 2024 12:46:46 -0700
Subject: [PATCH] Simplify stack management with _frame_state

---
 include/unifex/task.hpp | 59 ++++++++++++++++++++++++++---------------
 1 file changed, 38 insertions(+), 21 deletions(-)

diff --git a/include/unifex/task.hpp b/include/unifex/task.hpp
index 535e1ed1..a51fae42 100644
--- a/include/unifex/task.hpp
+++ b/include/unifex/task.hpp
@@ -416,6 +416,23 @@ struct _promise final {
   };
 };
 
+struct _frame_state {
+  _frame_state() noexcept = default;
+
+  explicit _frame_state(AsyncStackFrame& frame, AsyncStackRoot& root) noexcept
+    : frame_(&frame), root_(&root) {}
+
+  void restore_frame_state() const noexcept {
+    if (frame_) {
+      activateAsyncStackFrame(*root_, *frame_);
+    }
+  }
+
+private:
+  AsyncStackFrame* frame_{};
+  AsyncStackRoot* root_; // only conditionally initialized
+};
+
 struct _sr_thunk_promise_base : _promise_base {
   _sr_thunk_promise_base()
     : _promise_base([this]() noexcept -> coro::coroutine_handle<> {
@@ -423,19 +440,10 @@ struct _sr_thunk_promise_base : _promise_base {
 
       whoToContinue_ = continuation_.done_handle();
 
-      AsyncStackRoot* root;
-      auto* const frame = frame_;
-      if (frame) {
-        popAsyncStackFrameFromCaller(*frame);
-
-        root = frame_->getStackRoot();
-        deactivateAsyncStackFrame(*frame);
-      }
+      const auto frameState = ensure_frame_deactivated();
 
       if (refCount_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
-        if (frame) {
-          activateAsyncStackFrame(*root, *frame);
-        }
+        frameState.restore_frame_state();
 
         return whoToContinue_;
       } else {
@@ -518,6 +526,22 @@ struct _sr_thunk_promise_base : _promise_base {
   void register_stop_callback() noexcept {
     callback_.construct(stoken_, stop_callback{this});
   }
+
+  _frame_state ensure_frame_deactivated() noexcept {
+    if (frame_ != nullptr) {
+      if (whoToContinue_ == continuation_.done_handle()) {
+        popAsyncStackFrameFromCaller(*frame_);
+      }
+
+      auto* root = frame_->getStackRoot();
+      // this asserts that root is not null
+      deactivateAsyncStackFrame(*frame_);
+
+      return _frame_state(*frame_, *root);
+    }
+
+    return {};
+  }
 };
 
 // TODO: determine if this should also be nothrow
@@ -560,20 +584,13 @@ struct _sr_thunk_promise final {
           // deferred stop callback's completion
           p.whoToContinue_ = p.continuation_.handle();
 
-          AsyncStackRoot* root;
-          auto* const frame = p.frame_;
-
-          if (frame) {
-            root = frame->getStackRoot();
-            deactivateAsyncStackFrame(*frame);
-          }
+          const auto frameState = p.ensure_frame_deactivated();
 
           // if we're last to complete, continue our continuation; otherwise do
           // nothing and wait for the async stop request to do it
           if (p.refCount_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
-            if (frame) {
-              activateAsyncStackFrame(*root, *frame);
-            }
+            frameState.restore_frame_state();
+
             return p.whoToContinue_;
           } else {
             // the deferred stop callback will reactivate this frame