continuation.hpp

#ifndef MRS_LIB_CORO_INTERNAL_CONTINUATION_HPP_
#define MRS_LIB_CORO_INTERNAL_CONTINUATION_HPP_
 
 
#include <coroutine>
#include <functional>
#include <stop_token>
#include <utility>
 
namespace mrs_lib::coro::internal
{
 
  template <typename T>
  std::coroutine_handle<T> coroutine_handle_reinterpret_cast(std::coroutine_handle<> handle)
  {
    return std::coroutine_handle<T>::from_address(handle.address());
  }
 
  template <typename T>
  struct CancellableContinuationFor;
 
  template <>
  struct CancellableContinuationFor<void>
  {
  };
 
  class CancellableContinuation
  {
  private:
    template <typename T>
    using ContinuationGetterFptr = CancellableContinuation (*)(std::coroutine_handle<T>);
    template <typename T>
    using TokenGetterFptr = std::stop_token (*)(std::coroutine_handle<T>);
 
  public:
    CancellableContinuation() = default;
 
    explicit CancellableContinuation(std::coroutine_handle<void>) = delete;
 
    template <typename T>
    explicit CancellableContinuation(std::coroutine_handle<T> handle)
        : data_{
              .func = [](std::coroutine_handle<> handle) -> CancellableContinuation {
                using Trait = CancellableContinuationFor<T>;
                static_assert(std::same_as<decltype(&Trait::release_continuation), ContinuationGetterFptr<T>>,
                              "Wrong signature for CancellableContinuationFor<T>::release_continuation.");
                static_assert(std::same_as<decltype(&Trait::get_token), TokenGetterFptr<T>>, "Wrong signature for CancellableContinuationFor<T>::get_token.");
                return Trait::release_continuation(coroutine_handle_reinterpret_cast<T>(handle));
              },
              .handle = handle,
              .stop_token = CancellableContinuationFor<T>::get_token(coroutine_handle_reinterpret_cast<T>(handle)),
          }
    {
    }
 
    ~CancellableContinuation()
    {
      cancel_and_destroy();
    }
 
    CancellableContinuation(const CancellableContinuation&) = delete;
    CancellableContinuation& operator=(const CancellableContinuation&) = delete;
    CancellableContinuation(CancellableContinuation&& other) noexcept : data_(std::exchange(other.data_, {}))
    {
    }
    CancellableContinuation& operator=(CancellableContinuation&& other) noexcept
    {
      std::ranges::swap(data_, other.data_);
      return *this;
    }
 
    std::coroutine_handle<> release()
    {
      return std::exchange(data_.handle, nullptr);
    }
 
    void cancel_and_destroy()
    {
      ContinuationData local_data = release_data();
 
      while (local_data.handle != nullptr)
      {
        CancellableContinuation continuation = std::invoke(local_data.func, local_data.handle);
        local_data.handle.destroy();
        local_data = continuation.release_data();
      }
    }
 
    std::stop_token get_token() const
    {
      return data_.stop_token;
    }
 
    friend bool operator==(const CancellableContinuation& continuation, std::nullptr_t)
    {
      return continuation.data_.handle == nullptr;
    }
 
  private:
    struct ContinuationData
    {
      ContinuationGetterFptr<void> func = nullptr;
      std::coroutine_handle<> handle = nullptr;
      std::stop_token stop_token{};
    };
 
    ContinuationData release_data()
    {
      return std::exchange(data_, {});
    }
 
    ContinuationData data_;
  };
 
} // namespace mrs_lib::coro::internal
 
#endif // MRS_LIB_CORO_INTERNAL_CONTINUATION_HPP_