task.hpp

#ifndef MRS_LIB_CORO_TASK_HPP_
#define MRS_LIB_CORO_TASK_HPP_
 
#include <cassert>
#include <concepts>
#include <coroutine>
#include <functional>
#include <memory>
#include <type_traits>
#include <utility>
 
#include <mrs_lib/coro/internal/attributes.hpp>
#include <variant>
 
// Note on ownership semantics:
// Since we want to support cancellation at any point in the coroutine stacks,
// each coroutine is owned by the object responsible for resuming it. If the
// coroutine is running, it is responsible for it's own lifetime - it has to
// either suspend and become continuation of some other task thus transferring
// ownership or destruct itself once completed.
 
namespace mrs_lib
{
 
  template <typename T = void>
    requires(std::same_as<T, std::remove_cvref_t<T>>)
  class Task;
 
  namespace internal
  {
 
    template <typename T>
    struct CoroutineDestroyer
    {
      void operator()(std::coroutine_handle<T> handle)
      {
        handle.destroy();
      }
      using pointer = std::coroutine_handle<T>;
    };
 
    template <typename T = void>
    using OwningCoroutineHandle = std::unique_ptr<std::coroutine_handle<T>, CoroutineDestroyer<T>>;
 
    template <typename T>
    class DeferredCoroutineDestroyer
    {
    public:
      DeferredCoroutineDestroyer(std::coroutine_handle<T> handle) : handle_(handle)
      {
      }
      ~DeferredCoroutineDestroyer()
      {
        std::invoke(CoroutineDestroyer<T>{}, handle_);
      }
      DeferredCoroutineDestroyer(const DeferredCoroutineDestroyer&) = delete;
      DeferredCoroutineDestroyer& operator=(const DeferredCoroutineDestroyer&) = delete;
      DeferredCoroutineDestroyer(DeferredCoroutineDestroyer&&) = delete;
      DeferredCoroutineDestroyer& operator=(DeferredCoroutineDestroyer&&) = delete;
 
    private:
      std::coroutine_handle<T> handle_;
    };
 
    template <typename Derived>
    class BasePromiseType
    {
      class FinalAwaitable
      {
      public:
        FinalAwaitable() = default;
        ~FinalAwaitable() = default;
        FinalAwaitable(const FinalAwaitable&) = delete;
        FinalAwaitable& operator=(const FinalAwaitable&) = delete;
        FinalAwaitable(FinalAwaitable&&) = delete;
        FinalAwaitable& operator=(FinalAwaitable&&) = delete;
 
        // Always suspend the ending task
        bool await_ready() noexcept;
 
        // SYMMETRIC TRANSFER IS BROKEN IN GCC and can result in stack
        // overflow when many tasks complete synchronously.
        // Because of this problem, the `await_suspend` uses the void signature
        // and resumes the continuation on a thread-local scheduler as a workaround.
        void await_suspend(std::coroutine_handle<Derived> task_handle) noexcept;
 
        // This should be unreachable - ended task should not be resumed
        void await_resume() noexcept;
      };
 
    public:
      // The task is lazy and will only start when awaited
      std::suspend_always initial_suspend();
      // The coroutine will be suspended and the continuation will be resumed
      FinalAwaitable final_suspend() noexcept;
 
      void set_continuation(OwningCoroutineHandle<> continuation);
 
    private:
      OwningCoroutineHandle<> continuation_{std::noop_coroutine()};
    };
 
    template <typename T>
    class ResultStorage
    {
    private:
      // Not enum class to allow usage in functions like std::get
      enum State : size_t
      {
        empty = 0,
        value = 1,
        exception = 2,
      };
 
    public:
      constexpr ResultStorage() noexcept : data_()
      {
      }
 
      constexpr void set_value(T&& val) noexcept(std::is_nothrow_move_constructible_v<T>)
      {
        assert(data_.index() == State::empty);
        data_.template emplace<State::value>(std::move(val));
      }
 
      void set_exception(std::exception_ptr eptr) noexcept
      {
        assert(data_.index() == State::empty || data_.valueless_by_exception());
        data_.template emplace<State::exception>(std::move(eptr));
      }
 
      constexpr T get_value() &&
      {
        size_t state = data_.index();
        if (state == State::exception)
        {
          std::rethrow_exception(std::get<State::exception>(data_));
        }
        assert(state == State::value);
        return std::get<State::value>(std::move(data_));
      }
 
    private:
      std::variant<std::monostate, T, std::exception_ptr> data_;
    };
 
    template <typename T>
    class PromiseType : public BasePromiseType<PromiseType<T>>
    {
    public:
      Task<T> get_return_object();
 
      void return_value(T&& ret_val);
 
      void unhandled_exception();
 
      T get_value()
      {
        return std::move(result_).get_value();
      }
 
    private:
      ResultStorage<T> result_;
    };
 
    template <>
    class PromiseType<void> : public BasePromiseType<PromiseType<void>>
    {
    public:
      Task<void> get_return_object();
 
      void return_void();
 
      void unhandled_exception();
 
      void get_value()
      {
        if (exception_)
        {
          std::rethrow_exception(exception_);
        }
      }
 
    private:
      std::exception_ptr exception_;
    };
 
    template <typename T>
    class TaskAwaitable
    {
      using Promise = Task<T>::promise_type;
 
    public:
      bool await_ready()
      {
        return false;
      }
 
      std::coroutine_handle<> await_suspend(std::coroutine_handle<> continuation)
      {
        task_handle_.promise().set_continuation(OwningCoroutineHandle<>(continuation));
        return task_handle_;
      }
 
      T await_resume()
      {
        DeferredCoroutineDestroyer destroyer{this->task_handle_};
        return this->task_handle_.promise().get_value();
      }
 
      ~TaskAwaitable() = default;
      TaskAwaitable(const TaskAwaitable&) = delete;
      TaskAwaitable& operator=(const TaskAwaitable&) = delete;
      TaskAwaitable(TaskAwaitable&&) = delete;
      TaskAwaitable& operator=(TaskAwaitable&&) = delete;
 
    private:
      TaskAwaitable(std::coroutine_handle<Promise> task_handle) : task_handle_(task_handle)
      {
      }
 
      std::coroutine_handle<Promise> task_handle_;
 
      friend class Task<T>;
    };
 
  } // namespace internal
 
  template <typename T>
    requires(std::same_as<T, std::remove_cvref_t<T>>)
  class [[nodiscard("Task is lazy and does not run until `co_await`ed.")]] MRS_LIB_INTERNAL_CORO_RETURN_TYPE MRS_LIB_INTERNAL_CORO_LIFETIMEBOUND Task
  {
  public:
    using promise_type = internal::PromiseType<T>;
 
    ~Task() = default;
    Task(const Task&) = delete;
    Task& operator=(const Task&) = delete;
    Task(Task&&) = delete;
    Task& operator=(Task&&) = delete;
 
    friend internal::TaskAwaitable<T> operator co_await(Task task)
    {
      return internal::TaskAwaitable<T>(task.coroutine_.release());
    }
 
  private:
    Task(internal::OwningCoroutineHandle<promise_type> coroutine) : coroutine_(std::move(coroutine))
    {
    }
 
    internal::OwningCoroutineHandle<promise_type> coroutine_;
 
    friend class internal::PromiseType<T>;
  };
 
} // namespace mrs_lib
 
#ifndef MRS_LIB_CORO_TASK_IMPL_HPP_
#include <mrs_lib/coro/task.impl.hpp> // IWYU pragma: export
#endif
 
#endif // MRS_LIB_CORO_TASK_HPP_