promise.h

Go to the documentation of this file.

// uvco (c) 2023-2024 Lewin Bormann. See LICENSE for specific terms.
 
#pragma once
 
#include "uvco/exception.h"
#include "uvco/internal/internal_utils.h"
#include "uvco/promise/promise_core.h"
 
#include <boost/assert.hpp>
#include <fmt/format.h>
 
#include <coroutine>
#include <exception>
#include <utility>
 
namespace uvco {
 
 
template <typename T> class Coroutine;
template <typename T> class Promise;
 
template <typename T> class PromiseHandle {
public:
  PromiseHandle(const PromiseHandle &) = delete;
  PromiseHandle(PromiseHandle &&) = delete;
  PromiseHandle &operator=(const PromiseHandle &) = delete;
  PromiseHandle &operator=(PromiseHandle &&) = delete;
  ~PromiseHandle() {
    if (core_ != nullptr) {
      core_->delRef();
    }
  }
 
  void cancel() {
    if (core_ != nullptr) {
      core_->cancel();
    }
  }
 
private:
  explicit PromiseHandle(PromiseCore<T> *core) : core_{core->addRef()} {}
 
  friend class Promise<T>;
  PromiseCore<T> *core_;
};
 
template <typename T> class Promise {
protected:
  struct PromiseAwaiter_;
  using PromiseCore_ = PromiseCore<T>;
  using SharedCore_ = PromiseCore_ *;
 
public:
  using promise_type = Coroutine<T>;
 
  Promise() : core_{makeRefCounted<PromiseCore_>()} {}
  explicit Promise(T &&result)
      : core_{makeRefCounted<PromiseCore_>(std::move(result))} {}
 
  Promise(Promise<T> &&other) noexcept : core_{other.core_} {
    other.core_ = nullptr;
  }
  Promise &operator=(const Promise<T> &other) {
    if (this == &other) {
      return *this;
    }
    core_ = other.core_->addRef();
    return *this;
  }
  Promise &operator=(Promise<T> &&other) noexcept {
    if (this == &other) {
      return *this;
    }
    if (core_ != nullptr) {
      core_->delRef();
    }
    core_ = other.core_;
    other.core_ = nullptr;
    return *this;
  }
  // A promise can be copied at low cost.
  Promise(const Promise<T> &other) : core_{other.core_->addRef()} {}
  ~Promise() {
    if (core_ != nullptr) {
      core_->delRef();
    }
  }
 
  PromiseHandle<T> handle() { return PromiseHandle<T>{core_}; }
 
  PromiseAwaiter_ operator co_await() const { return PromiseAwaiter_{*core_}; }
 
  [[nodiscard]] bool ready() const { return core_->ready(); }
 
  T unwrap() {
    if (ready()) {
      auto &slot = core_->slot.value();
      switch (slot.index()) {
      case 0: {
        T value = std::move(std::get<0>(slot));
        core_->slot.reset();
        return std::move(value);
      }
      case 1: {
        std::rethrow_exception(std::get<1>(slot));
      }
      default:
        throw UvcoException("PromiseAwaiter_::await_resume: invalid slot");
      }
    }
    throw UvcoException("unwrap called on unfulfilled promise");
  }
 
protected:
  struct PromiseAwaiter_ {
    explicit PromiseAwaiter_(PromiseCore_ &core) : core_{core} {}
    PromiseAwaiter_(PromiseAwaiter_ &&) = delete;
    PromiseAwaiter_(const PromiseAwaiter_ &) = delete;
    PromiseAwaiter_ &operator=(PromiseAwaiter_ &&) = delete;
    PromiseAwaiter_ &operator=(const PromiseAwaiter_ &) = delete;
    ~PromiseAwaiter_() = default;
 
    [[nodiscard]] bool await_ready() const { return core_.ready(); }
    [[nodiscard]] bool await_suspend(std::coroutine_handle<> handle) const {
      BOOST_ASSERT_MSG(!core_.willResume(),
                       "promise is already being waited on!");
      core_.setHandle(handle);
      return true;
    }
    T await_resume() const {
      if (core_.stale()) {
        throw UvcoException(
            "co_await called on previously finished promise (T)");
      }
      if (core_.slot.has_value()) {
        switch (core_.slot->index()) {
        case 0: {
          T result = std::move(std::get<0>(core_.slot.value()));
          core_.slot.reset();
          return std::move(result);
        }
        case 1:
          std::rethrow_exception(std::get<1>(core_.slot.value()));
        default:
          throw UvcoException("PromiseAwaiter_::await_resume: invalid slot");
        }
      }
      throw UvcoException("await_resume called on unfulfilled promise (bug?)");
    }
 
    PromiseCore_ &core_;
  };
 
  template <typename U> friend class Coroutine;
  template <typename... Ts> friend class SelectSet;
 
  explicit Promise(SharedCore_ core) : core_{core->addRef()} {}
  SharedCore_ &core() { return core_; }
 
  SharedCore_ core_;
};
 
template <> class Promise<void> {
  struct PromiseAwaiter_;
  using SharedCore_ = PromiseCore<void> *;
 
public:
  using promise_type = Coroutine<void>;
 
  Promise();
  Promise(Promise<void> &&other) noexcept;
  Promise &operator=(const Promise<void> &other);
  Promise &operator=(Promise<void> &&other) noexcept;
  Promise(const Promise<void> &other);
  ~Promise();
 
  PromiseHandle<void> handle();
 
  PromiseAwaiter_ operator co_await() const;
 
  [[nodiscard]] bool ready() const;
 
  // Get the result *right now*, and throw an exception if the promise
  // is not ready, or if it encountered an exception itself.
  void unwrap();
 
private:
  struct PromiseAwaiter_ {
    explicit PromiseAwaiter_(PromiseCore<void> &core);
    PromiseAwaiter_(PromiseAwaiter_ &&) = delete;
    PromiseAwaiter_(const PromiseAwaiter_ &) = delete;
    PromiseAwaiter_ &operator=(PromiseAwaiter_ &&) = delete;
    PromiseAwaiter_ &operator=(const PromiseAwaiter_ &) = delete;
 
    [[nodiscard]] bool await_ready() const;
    [[nodiscard]] bool await_suspend(std::coroutine_handle<> handle) const;
    void await_resume() const;
 
    PromiseCore<void> &core_;
  };
 
  explicit Promise(SharedCore_ core);
 
  friend class Coroutine<void>;
  template <typename... Ts> friend class SelectSet;
 
  SharedCore_ &core() { return core_; }
 
  SharedCore_ core_;
};
 
template <typename T> class Coroutine {
  using PromiseCore_ = PromiseCore<T>;
  using SharedCore_ = PromiseCore_ *;
 
public:
  Coroutine() : core_{makeRefCounted<PromiseCore_>()} {}
  Coroutine(const Coroutine &other) = delete;
  Coroutine &operator=(const Coroutine &other) = delete;
  Coroutine(Coroutine &&other) = delete;
  Coroutine &operator=(Coroutine &&other) = delete;
 
  ~Coroutine() {
    if (core_ != nullptr) {
      core_->delRef();
    }
  }
 
  Promise<T> get_return_object() { return Promise<T>{core_}; }
 
  void return_value(T value) {
    // Probably cancelled.
    if (core_->slot.has_value() && core_->slot->index() == 1) {
      return;
    }
    BOOST_ASSERT(!core_->slot);
    core_->slot = std::move(value);
    core_->resume();
  }
 
  // Note: if suspend_always is chosen, we can better control when the promise
  // will be scheduled.
  std::suspend_never initial_suspend() noexcept { return {}; }
  std::suspend_never final_suspend() noexcept { return {}; }
 
  // Part of the coroutine protocol: called upon unhandled exception leaving the
  // coroutine.
  void unhandled_exception() {
    core_->except(std::current_exception());
    core_->resume();
  }
 
protected:
  SharedCore_ core_;
};
 
template <> class Coroutine<void> {
  using PromiseCore_ = PromiseCore<void>;
  using SharedCore_ = PromiseCore_ *;
 
public:
  Coroutine() : core_{makeRefCounted<PromiseCore_>()} {}
  // Coroutine is pinned in memory and not allowed to copy/move.
  Coroutine(Coroutine<void> &&other) noexcept = delete;
  Coroutine &operator=(const Coroutine<void> &other) = delete;
  Coroutine &operator=(Coroutine<void> &&other) = delete;
  Coroutine(const Coroutine<void> &other) = delete;
  ~Coroutine() {
    if (core_ != nullptr) {
      core_->delRef();
    }
  }
 
  Promise<void> get_return_object() { return Promise<void>{core_}; }
  std::suspend_never initial_suspend() noexcept { return {}; }
  std::suspend_never final_suspend() noexcept { return {}; }
 
  void return_void();
  void unhandled_exception();
 
private:
  SharedCore_ core_;
};
 
 
} // namespace uvco