Better support linking TUs with mixed async stacks

After this change, we work harder at supporting linking together TUs that differ in whether they support async stacks or not.
facebookexperimental · Sep 3, 2024 · 258ef37 · 258ef37
1 parent 297be31
commit 258ef37
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 41 deletions.
diff --git a/include/unifex/await_transform.hpp b/include/unifex/await_transform.hpp
@@ -97,7 +97,9 @@ struct _awaitable_base<Promise, Value, WithAsyncStackSupport>::type {
         if (auto* frame = get_async_stack_frame(continuation_.promise())) {
           detail::ScopedAsyncStackRoot root;
           root.activateFrame(*frame);
-          return continuation_.resume();
+          continuation_.resume();
+          root.ensureFrameDeactivated(frame);
+          return;
         }
       }
 
@@ -142,7 +144,9 @@ struct _awaitable_base<Promise, Value, WithAsyncStackSupport>::type {
           detail::ScopedAsyncStackRoot root;
           root.activateFrame(frame);
 
-          return continuation_.resume_done();
+          continuation_.resume_done();
+          root.ensureFrameDeactivated(&frame);
+          return;
         }
       }
 
@@ -213,7 +217,7 @@ struct _awaitable<Promise, Sender, WithAsyncStackSupport>::type
   void await_suspend(coro::coroutine_handle<Promise> handle) noexcept {
     if constexpr (WithAsyncStackSupport) {
       auto* frame = get_async_stack_frame(handle.promise());
-      if (frame) {
+      if (frame && frame->getStackRoot()) {
         deactivateAsyncStackFrame((*frame));
       }
     }

diff --git a/include/unifex/sender_concepts.hpp b/include/unifex/sender_concepts.hpp
@@ -250,6 +250,7 @@ namespace _cpo {
 
 struct _fn {
 private:
+  template <bool WithAsyncStackSupport>
   struct _impl {
     template(typename S, typename R)         //
         (requires tag_invocable<_fn, S, R>)  //
@@ -271,31 +272,38 @@ struct _fn {
     }
   };
 
-#if UNIFEX_NO_ASYNC_STACKS
-public:
-  template(typename S, typename R)          //
-      (requires sender<S> AND receiver<R>)  //
-      auto
-      operator()(S&& s, R&& r) const
-      noexcept(noexcept(_impl{}(std::forward<S>(s), std::forward<R>(r))))
-          -> decltype(_impl{}(std::forward<S>(s), std::forward<R>(r))) {
-    return _impl{}(std::forward<S>(s), std::forward<R>(r));
-  }
-#else
   template <typename S, typename R>
-  using op_t = _inject::
-      op_wrapper<std::invoke_result_t<_impl, S, _inject::receiver_t<R>>, R>;
+  using op_t = _inject::op_wrapper<
+      std::invoke_result_t<_impl<false>, S, _inject::receiver_t<R>>,
+      R>;
+
+  template <>
+  struct _impl<true> {
+    template <typename S, typename R>
+    auto operator()(S&& s, R&& r) const
+        noexcept(noexcept(_inject::make_op_wrapper(
+            std::forward<S>(s), std::forward<R>(r), _impl<false>{})))
+            -> op_t<S, R> {
+      return _inject::make_op_wrapper(
+          std::forward<S>(s), std::forward<R>(r), _impl<false>{});
+    }
+  };
 
 public:
-  template(typename S, typename R)          //
-      (requires sender<S> AND receiver<R>)  //
+  template(
+      typename S,
+      typename R,
+      bool WithAsyncStackSupport = !UNIFEX_NO_ASYNC_STACKS)  //
+      (requires sender<S> AND receiver<R>)                   //
       auto
-      operator()(S&& s, R&& r) const noexcept(noexcept(_inject::make_op_wrapper(
-          std::forward<S>(s), std::forward<R>(r), _impl{}))) -> op_t<S, R> {
-    return _inject::make_op_wrapper(
-        std::forward<S>(s), std::forward<R>(r), _impl{});
+      operator()(S&& s, R&& r) const
+      noexcept(noexcept(_impl<WithAsyncStackSupport>{}(
+          std::forward<S>(s), std::forward<R>(r))))
+          -> decltype(_impl<WithAsyncStackSupport>{}(
+              std::forward<S>(s), std::forward<R>(r))) {
+    return _impl<WithAsyncStackSupport>{}(
+        std::forward<S>(s), std::forward<R>(r));
   }
-#endif
 };
 
 }  // namespace _cpo

diff --git a/include/unifex/tracing/inject_async_stack.hpp b/include/unifex/tracing/inject_async_stack.hpp
@@ -17,21 +17,19 @@
 
 #include <unifex/config.hpp>
 
-#if !UNIFEX_NO_ASYNC_STACKS
+#include <unifex/continuations.hpp>
+#include <unifex/receiver_concepts.hpp>
+#include <unifex/tracing/async_stack.hpp>
+#include <unifex/tracing/get_async_stack_frame.hpp>
+#include <unifex/tracing/get_return_address.hpp>
+#include <unifex/detail/unifex_fwd.hpp>
 
-#  include <unifex/continuations.hpp>
-#  include <unifex/receiver_concepts.hpp>
-#  include <unifex/tracing/async_stack.hpp>
-#  include <unifex/tracing/get_async_stack_frame.hpp>
-#  include <unifex/tracing/get_return_address.hpp>
-#  include <unifex/detail/unifex_fwd.hpp>
+#include <exception>
+#include <functional>
+#include <type_traits>
+#include <utility>
 
-#  include <exception>
-#  include <functional>
-#  include <type_traits>
-#  include <utility>
-
-#  include <unifex/detail/prologue.hpp>
+#include <unifex/detail/prologue.hpp>
 
 namespace unifex {
 namespace _inject {
@@ -172,13 +170,13 @@ struct _rcvr_wrapper<Receiver>::type final : _rcvr_base {
     return std::move(cpo)(std::as_const(r.receiver()));
   }
 
-#  if UNIFEX_ENABLE_CONTINUATION_VISITATIONS
+#if UNIFEX_ENABLE_CONTINUATION_VISITATIONS
   template <typename Visit>
   friend void
   tag_invoke(tag_t<visit_continuations>, const type& r, Visit&& visit) {
     std::invoke(visit, r.receiver());
   }
-#  endif
+#endif
 
   Receiver& receiver() const noexcept {
     return static_cast<op_with_receiver_t<Receiver>*>(this->op_)->receiver_;
@@ -238,6 +236,4 @@ auto make_op_wrapper(S&& s, R&& r, Fn&& fn) noexcept(
 
 }  // namespace unifex
 
-#  include <unifex/detail/epilogue.hpp>
-
-#endif  // !UNIFEX_NO_ASYNC_STACKS
+#include <unifex/detail/epilogue.hpp>