Refactor fork reclamation strategies into policies

rt/include/typon/fork.hpp

@@ -5,6 +5,7 @@
 #include <coroutine>
 #include <cstdint>
 
+#include <typon/defer.hpp>
 #include <typon/forked.hpp>
 #include <typon/result.hpp>
 #include <typon/scheduler.hpp>
@@ -14,7 +15,73 @@
 namespace typon
 {
 
-  template <typename T = void>
+  namespace policy
+  {
+    struct Bundle
+    {
+      void on_final_suspend(std::coroutine_handle<> coroutine) noexcept
+      {
+        (void) coroutine;
+      }
+
+      struct OnAwaitable
+      {
+        template <typename Promise>
+        void on_await_suspend(std::coroutine_handle<Promise> coroutine) noexcept
+        {
+          (void) coroutine;
+        }
+
+        template <typename Promise>
+        auto on_await_resume(std::coroutine_handle<Promise> coroutine)
+        {
+          using T = typename Promise::value_type;
+          auto thefts = coroutine.promise()._span->_thefts;
+          auto rank = coroutine.promise()._rank;
+          bool ready = (thefts == rank);
+          if (!ready)
+          {
+            coroutine.promise()._span->_children.push_back(coroutine);
+          }
+          return Forked<T>(coroutine, ready, nullptr);
+        }
+      };
+    };
+
+    struct Refcnt
+    {
+      ForkNode _node;
+
+      void on_final_suspend(std::coroutine_handle<> coroutine) noexcept
+      {
+        (void) coroutine;
+        _node.decref();
+      }
+
+      struct OnAwaitable
+      {
+        template <typename Promise>
+        void on_await_suspend(std::coroutine_handle<Promise> coroutine) noexcept
+        {
+          coroutine.promise()._policy._node._coroutine = coroutine;
+        }
+
+        template <typename Promise>
+        auto on_await_resume(std::coroutine_handle<Promise> coroutine)
+        {
+          using T = typename Promise::value_type;
+          auto thefts = coroutine.promise()._span->_thefts;
+          auto rank = coroutine.promise()._rank;
+          bool ready = (thefts == rank);
+          auto node = &(coroutine.promise()._policy._node);
+          return Forked<T>(coroutine, ready, node);
+        }
+      };
+    };
+
+  }
+
+  template <typename T = void, typename Policy = policy::Refcnt>
   struct [[nodiscard]] Fork
   {
     struct promise_type;
@@ -36,11 +103,7 @@ namespace typon
     {
       Span * _span;
       u64 _rank;
-      ForkNode _node;
-
-      promise_type() noexcept
-        : _node{ std::coroutine_handle<promise_type>::from_promise(*this) }
-      {}
+      [[no_unique_address]] Policy _policy;
 
       Fork get_return_object() noexcept
       {
@@ -63,19 +126,11 @@ namespace typon
             {
               return span->_coroutine;
             }
-            auto rank = coroutine.promise()._rank;
             if (auto & exception = coroutine.promise()._exception)
             {
-              span->set_exception(exception, rank);
-            }
-            if (!(rank & 1))
-            {
-              auto & ref = coroutine.promise()._node._ref;
-              if (!ref.exchange(false, std::memory_order_acq_rel))
-              {
-                coroutine.destroy();
-              }
+              span->set_exception(exception, coroutine.promise()._rank);
             }
+            coroutine.promise()._policy.on_final_suspend(coroutine);
             u64 n = span->_n.fetch_sub(1, std::memory_order_acq_rel);
             if (n == 1)
             {
@@ -92,42 +147,16 @@ namespace typon
     struct awaitable : std::suspend_always
     {
       std::coroutine_handle<promise_type> _coroutine;
+      [[no_unique_address]] Policy::OnAwaitable _policy;
 
       template <typename Promise>
       auto await_suspend(std::coroutine_handle<Promise> continuation) noexcept
       {
         Span * span = &(continuation.promise()._span);
         _coroutine.promise()._span = span;
-        _coroutine.promise()._rank = (span->_thefts << 1);
+        _coroutine.promise()._rank = span->_thefts;
 
-        std::coroutine_handle<> on_stack_handle = _coroutine;
-        Scheduler::push(span);
-        return on_stack_handle;
-      }
-
-      auto await_resume()
-      {
-        auto thefts = _coroutine.promise()._span->_thefts;
-        auto rank = _coroutine.promise()._rank;
-        return Forked<T>(_coroutine, (thefts == (rank >> 1)), true);
-      }
-    };
-
-    auto operator co_await() &&
-    {
-      return awaitable { {}, _coroutine };
-    }
-
-    struct noloop_awaitable : std::suspend_always
-    {
-      std::coroutine_handle<promise_type> _coroutine;
-
-      template <typename Promise>
-      auto await_suspend(std::coroutine_handle<Promise> continuation) noexcept
-      {
-        Span * span = &(continuation.promise()._span);
-        _coroutine.promise()._span = span;
-        _coroutine.promise()._rank = (span->_thefts << 1) + 1;
+        _policy.on_await_suspend(_coroutine);
 
         std::coroutine_handle<> on_stack_handle = _coroutine;
         Scheduler::push(span);
@@ -136,20 +165,13 @@ namespace typon
 
       auto await_resume()
       {
-        auto thefts = _coroutine.promise()._span->_thefts;
-        auto rank = _coroutine.promise()._rank;
-        bool ready = thefts == (rank >> 1);
-        if (!ready)
-        {
-          _coroutine.promise()._span->_children.push_back(_coroutine);
-        }
-        return Forked<T>(_coroutine, ready, false);
+        return _policy.on_await_resume(_coroutine);
       }
     };
 
-    auto noloop() &&
+    auto operator co_await() &&
     {
-      return noloop_awaitable { {}, _coroutine };
+      return awaitable { {}, _coroutine, {} };
     }
   };
 

rt/include/typon/forked.hpp

@@ -17,6 +17,14 @@ namespace typon
   {
     std::coroutine_handle<> _coroutine;
     std::atomic<bool> _ref {true};
+
+    void decref() noexcept
+    {
+      if (!_ref.exchange(false, std::memory_order_acq_rel))
+      {
+        _coroutine.destroy();
+      }
+    }
   };
 
 
@@ -110,8 +118,8 @@ namespace typon
 
     Result<T> * _result = nullptr;
 
-    template <typename Promise>
-    Forked(std::coroutine_handle<Promise> coroutine, bool ready, bool owning)
+    template <typename Coroutine>
+    Forked(Coroutine coroutine, bool ready, ForkNode * node)
     {
       if (ready)
       {
@@ -120,7 +128,7 @@ namespace typon
       }
       else
       {
-        this->_node = owning ? &(coroutine.promise()._node) : nullptr;
+        this->_node = node;
         _result = &(coroutine.promise());
       }
     }
@@ -163,10 +171,7 @@ namespace typon
       {
         if (auto node = this->_node)
         {
-          if (!node->_ref.exchange(false, std::memory_order_acq_rel))
-          {
-            node->_coroutine.destroy();
-          }
+          node->decref();
         }
       }
       else
@@ -194,20 +199,21 @@ namespace typon
 
     bool _ready;
     Result<T> * _result;
-    ForkNode * _node;
+    void * _coroutine;
 
-    template <typename Promise>
-    Forked(std::coroutine_handle<Promise> coroutine, bool ready, bool owning)
+    template <typename Coroutine>
+    Forked(Coroutine coroutine, bool ready, ForkNode * node)
     {
       _ready = ready;
       _result = &(coroutine.promise());
-      _node = (owning | ready) ? &(coroutine.promise()._node) : nullptr;
       if (ready)
       {
-        if (auto & exception = coroutine.promise()._exception)
-        {
-          std::rethrow_exception(exception);
+        _coroutine = coroutine.address();
+        coroutine.promise().get();
       }
+      else
+      {
+        _coroutine = node;
       }
     }
 
@@ -217,31 +223,28 @@ namespace typon
     Forked(Forked && other) noexcept
       : _ready(other._ready)
       , _result(other._result)
-      , _node(std::exchange(other._node, nullptr))
+      , _coroutine(std::exchange(other._coroutine, nullptr))
     {}
 
     Forked& operator=(Forked && other) noexcept
     {
       std::swap(_ready, other._ready);
-      std::swap(_node, other._node);
       std::swap(_result, other._result);
+      std::swap(_coroutine, other._coroutine);
       return *this;
     }
 
     ~Forked()
     {
-      if (_node)
+      if (_coroutine)
       {
         if (_ready)
         {
-          _node->_coroutine.destroy();
+          std::coroutine_handle<void>::from_address(_coroutine).destroy();
         }
         else
         {
-          if (!_node->_ref.exchange(false, std::memory_order_acq_rel))
-          {
-            _node->_coroutine.destroy();
-          }
+          reinterpret_cast<ForkNode *>(_coroutine)->decref();
         }
       }
     }

rt/include/typon/typon.hpp

@@ -26,6 +26,16 @@ namespace typon
   }
 
 
+  template <typename Policy, typename Task>
+  Fork<typename Task::promise_type::value_type, Policy> fork(Task task)
+  {
+    // Put the task in a local variable to ensure its destructor will
+    // be called on co_return instead of only on coroutine destruction.
+    Task local_task = std::move(task);
+    co_return co_await std::move(local_task);
+  }
+
+
   template <typename Task>
   Future<typename Task::promise_type::value_type> future(Task task)
   {