multipromise.h

Go to the documentation of this file.

 
#pragma once
 
#include "uvco/exception.h"
#include "uvco/loop/loop.h"
#include "uvco/promise/promise.h"
#include "uvco/promise/promise_core.h"
 
#include <coroutine>
#include <exception>
#include <memory>
#include <optional>
 
namespace uvco {
 
 
template <typename T> class MultiPromiseCore : public PromiseCore<T> {
public:
  MultiPromiseCore() = default;
  MultiPromiseCore(const MultiPromiseCore &) = delete;
  MultiPromiseCore(MultiPromiseCore &&) = delete;
  MultiPromiseCore &operator=(const MultiPromiseCore &) = delete;
  MultiPromiseCore &operator=(MultiPromiseCore &&) = delete;
  static_assert(!std::is_void_v<T>);
  ~MultiPromiseCore() override = default;
 
  void setHandle(std::coroutine_handle<> handle) override {
    // Once an external scheduler works, Promises will not be nested anymore
    // (resume called by resume down in the stack)
    //
    // BOOST_ASSERT(PromiseCore<T>::state_
    // == PromiseState::init || PromiseCore<T>::state_ ==
    // PromiseState::finished);
    //
    // state is init or running (latter can occur if setHandle is called from a
    // stack originating at resume()).
    BOOST_ASSERT_MSG(PromiseCore<T>::state_ != PromiseState::waitedOn,
                     "MultiPromise must be co_awaited before next yield");
    BOOST_ASSERT_MSG(!PromiseCore<T>::handle_,
                     "MultiPromise must be co_awaited before next yield");
    PromiseCore<T>::handle_ = handle;
    PromiseCore<T>::state_ = PromiseState::waitedOn;
  }
 
  void resume() override { PromiseCore<T>::resume(); }
 
  void resumeGenerator() {
    BOOST_ASSERT(PromiseCore<T>::state_ != PromiseState::finished);
    if (generatorHandle_) {
      const auto generatorHandle = generatorHandle_.value();
      generatorHandle_.reset();
      Loop::enqueue(generatorHandle);
    }
  }
 
  void suspendGenerator(std::coroutine_handle<> handle) {
    BOOST_ASSERT_MSG(!generatorHandle_,
                     "MultiPromiseCore::suspendGenerator: generatorHandle_ is "
                     "already set");
    generatorHandle_ = handle;
  }
 
  void cancelGenerator() {
    terminated();
    if (generatorHandle_) {
      const std::coroutine_handle<> generatorHandle = generatorHandle_.value();
      generatorHandle_.reset();
      // Careful: within this function, this class' dtor is called!
      generatorHandle.destroy();
    }
  }
 
  void terminated() { terminated_ = true; }
 
  [[nodiscard]] bool isTerminated() const { return terminated_; }
 
private:
  std::optional<std::coroutine_handle<>> generatorHandle_;
  bool terminated_ = false;
};
 
template <typename T> class Generator;
 
template <typename T> class MultiPromise {
protected:
  struct MultiPromiseAwaiter_;
 
  using PromiseCore_ = MultiPromiseCore<T>;
  using SharedCore_ = std::shared_ptr<PromiseCore_>;
 
public:
  using promise_type = Generator<T>;
  static_assert(!std::is_void_v<T>);
 
  MultiPromise(MultiPromise<T> &&) noexcept = default;
  MultiPromise &operator=(const MultiPromise<T> &) = delete;
  MultiPromise &operator=(MultiPromise<T> &&) noexcept = delete;
  MultiPromise(const MultiPromise<T> &other) = default;
  ~MultiPromise() {
    // Us and the coroutine frame; but we're about to be destroyed
    if (core_.use_count() == 2) {
      // -> cancel generator. Because the coroutine frame keeps a reference to
      // the core, we can't do this in the core's destructor.
      PromiseCore_ *weakCore = core_.get();
      core_.reset();
      // The core's dtor is called while cancelGenerator runs.
      weakCore->cancelGenerator();
    }
  }
 
  Promise<std::optional<T>> next() { co_return (co_await *this); }
 
  MultiPromiseAwaiter_ operator co_await() const {
    BOOST_ASSERT(core_);
    return MultiPromiseAwaiter_{core_};
  }
 
  bool ready() { return core_->slot.has_value() || core_->isTerminated(); }
 
  void cancel() { core_->cancelGenerator(); }
 
protected:
  struct MultiPromiseAwaiter_ {
    constexpr explicit MultiPromiseAwaiter_(SharedCore_ core)
        : core_{std::move(core)} {}
    MultiPromiseAwaiter_(MultiPromiseAwaiter_ &&) = delete;
    MultiPromiseAwaiter_(const MultiPromiseAwaiter_ &) = delete;
    MultiPromiseAwaiter_ &operator=(MultiPromiseAwaiter_ &&) = delete;
    MultiPromiseAwaiter_ &operator=(const MultiPromiseAwaiter_ &) = delete;
    ~MultiPromiseAwaiter_() = default;
 
    [[nodiscard]] bool await_ready() const {
      return core_->isTerminated() || core_->slot.has_value();
    }
    [[nodiscard]] bool await_suspend(std::coroutine_handle<> handle) const {
      BOOST_ASSERT_MSG(!core_->willResume(),
                       "promise is already being waited on!\n");
      core_->setHandle(handle);
      core_->resumeGenerator();
      return true;
    }
    std::optional<T> await_resume() const {
      if (!core_->slot) {
        // Terminated by co_return
        return std::nullopt;
      }
      switch (core_->slot->index()) {
      [[likely]] case 0: {
        std::optional<T> result = std::move(std::get<0>(core_->slot.value()));
        core_->slot.reset();
        return std::move(result);
      }
      case 1:
        // Terminated by exception
        BOOST_ASSERT(core_->isTerminated());
        std::rethrow_exception(std::get<1>(core_->slot.value()));
      default:
        throw UvcoException("MultiPromiseAwaiter_::await_resume: invalid slot");
      }
    }
 
    SharedCore_ core_;
  };
 
  template <typename U> friend class Generator;
 
  explicit MultiPromise(SharedCore_ core) : core_{std::move(core)} {}
 
  SharedCore_ core_;
};
 
template <typename T> class Generator {
  struct YieldAwaiter_;
  using PromiseCore_ = MultiPromiseCore<T>;
  using SharedCore_ = std::shared_ptr<PromiseCore_>;
 
public:
  // A generator promise object is pinned in memory (copy/move forbidden).
  Generator() : core_{std::make_shared<PromiseCore_>()} {}
  Generator(Generator<T> &&) noexcept = delete;
  Generator &operator=(const Generator<T> &) = delete;
  Generator &operator=(Generator<T> &&) noexcept = delete;
  Generator(const Generator<T> &other) = delete;
  ~Generator() = default;
 
  MultiPromise<T> get_return_object() { return MultiPromise<T>{core_}; }
 
  void return_void() {
    core_->terminated();
    core_->resume();
  }
 
  // Note: if suspend_always is chosen, we can better control when the
  // MultiPromise will be scheduled.
  std::suspend_never initial_suspend() noexcept { return {}; }
  std::suspend_never final_suspend() noexcept { return {}; }
 
  void unhandled_exception() {
    core_->slot = std::current_exception();
    core_->terminated();
    core_->resume();
  }
 
  YieldAwaiter_ yield_value(T value) {
    BOOST_ASSERT(!core_->slot);
    core_->slot = std::move(value);
    core_->resume();
    return YieldAwaiter_{*core_};
  }
 
private:
  struct YieldAwaiter_ {
    explicit YieldAwaiter_(PromiseCore_ &core) : core_{core} {}
 
    [[nodiscard]] bool await_ready() const { return !core_.slot.has_value(); }
 
    bool await_suspend(std::coroutine_handle<> handle) {
      core_.suspendGenerator(handle);
      return true;
    }
 
    void await_resume() {
      // Returning into the generator coroutine
    }
 
    PromiseCore_ &core_;
  };
 
  SharedCore_ core_;
};
 
 
} // namespace uvco