multipromise.h

Go to the documentation of this file.

// uvco (c) 2024 Lewin Bormann. See LICENSE for specific terms.
 
#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> {
  using PromiseCore<T>::coroutine_;
 
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>::waitingHandle_,
                     "MultiPromise must be co_awaited before next yield");
    PromiseCore<T>::waitingHandle_ = handle;
    PromiseCore<T>::state_ = PromiseState::waitedOn;
  }

  void resumeGenerator() {
    BOOST_ASSERT(PromiseCore<T>::state_ != PromiseState::finished);
    if (suspended_) {
      BOOST_ASSERT(coroutine_ != nullptr);
      suspended_ = false;
      Loop::enqueue(coroutine_);
    }
  }

  void suspendGenerator() {
    BOOST_ASSERT(!suspended_);
    suspended_ = true;
  }

  void cancelGenerator() {
    setTerminated();
    if (coroutine_ != nullptr) {
      const std::coroutine_handle<> coroutine = coroutine_;
      coroutine_ = nullptr;
      Loop::cancel(coroutine);
      // Careful: within this function, this class' dtor is called!
      coroutine.destroy();
    }
  }

  void setTerminated() { terminated_ = true; }

  [[nodiscard]] bool isTerminated() const { return terminated_; }
 
private:
  bool suspended_ = false;
  bool terminated_ = false;
};
 
template <typename T> class Generator;

template <typename T> class MultiPromise {
protected:
  struct MultiPromiseAwaiter_;
 
  using PromiseCore_ = MultiPromiseCore<T>;
 
public:
  using promise_type = Generator<T>;
  static_assert(!std::is_void_v<T>);

  MultiPromise(MultiPromise<T> &&other) noexcept : core_{other.core_} {
    other.core_ = nullptr;
  }
  MultiPromise &operator=(const MultiPromise<T> &) = delete;
  MultiPromise &operator=(MultiPromise<T> &&other) noexcept {
    if (this == &other) {
      return *this;
    }
    if (core_ != nullptr) {
      core_->destroyCoroutine();
    }
    core_ = other.core_;
    other.core_ = nullptr;
    return *this;
  }
  MultiPromise(const MultiPromise<T> &other) = delete;
  ~MultiPromise() {
    if (core_ != nullptr) {
      core_->cancelGenerator();
    }
    core_ = nullptr;
  }

  Promise<std::optional<T>> next() { co_return (co_await *this); }

  MultiPromiseAwaiter_ operator co_await() const {
    return MultiPromiseAwaiter_{core_};
  }

  bool ready() {
    return core_ == nullptr || core_->slot.has_value() || core_->isTerminated();
  }

  void cancel() {
    if (core_ == nullptr) {
      return;
    }
    core_->cancelGenerator();
    core_ = nullptr;
  }
 
protected:
  struct MultiPromiseAwaiter_ {
    constexpr explicit MultiPromiseAwaiter_(PromiseCore_ *core) : core_{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_ == nullptr || core_->isTerminated() ||
             core_->slot.has_value();
    }

    [[nodiscard]] std::coroutine_handle<>
    await_suspend(std::coroutine_handle<> handle) const {
      if (core_ == nullptr) {
        return Loop::getNext();
      }
      BOOST_ASSERT_MSG(!core_->isAwaited(),
                       "promise is already being waited on!\n");
      core_->setHandle(handle);
      core_->resumeGenerator();
      return Loop::getNext();
    }

    std::optional<T> await_resume() const {
      if (core_ == nullptr) {
        return std::nullopt;
      }
      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");
      }
    }
 
    PromiseCore_ *core_;
  };
 
  template <typename U> friend class Generator;
 
  explicit MultiPromise(PromiseCore_ &core) : core_{&core} {}
 
  PromiseCore_ *core_;
};

template <typename T> class Generator {
  struct YieldAwaiter_;
  using PromiseCore_ = MultiPromiseCore<T>;
 
public:
  // A generator promise object is pinned in memory (copy/move forbidden).
  Generator() = default;
  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_.setTerminated();
    core_.resume();
  }

  // Note: if suspend_always is chosen, we can better control when the
  // MultiPromise will be scheduled.
  std::suspend_never initial_suspend() noexcept {
    core_.setRunning(std::coroutine_handle<Generator<T>>::from_promise(*this));
    return {};
  }

  std::suspend_always final_suspend() noexcept { return {}; }

  void unhandled_exception() {
    core_.slot = std::current_exception();
    core_.setTerminated();
    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 */) {
      // The core_ already knows our handle from initial_suspend().
      core_.suspendGenerator();
      return true;
    }
 
    void await_resume() {
      // Returning into the generator coroutine
    }
 
    PromiseCore_ &core_;
  };
 
  PromiseCore_ core_;
};

 
} // namespace uvco