corochain.hpp

Go to the documentation of this file.

#pragma once
 
#include <coroutine>
#include <optional>
#include <expected>
#include <memory>
#include <functional>
#include <exception>
#include <utility>
 
#include "promises.hpp"
#include "poller.hpp"
 
namespace NNet {
 
template<typename T> struct TFinalAwaiter;
 
template<typename T> struct TFuture;
 
template<typename T>
struct TPromiseBase {
    std::suspend_never initial_suspend() { return {}; }
    TFinalAwaiter<T> final_suspend() noexcept;
    std::coroutine_handle<> Caller = std::noop_coroutine();
};
 
template<typename T>
struct TPromise: public TPromiseBase<T> {
    TFuture<T> get_return_object();
 
    void return_value(const T& t) {
        ErrorOr = t;
    }
 
    void return_value(T&& t) {
        ErrorOr = std::move(t);
    }
 
    void unhandled_exception() {
        ErrorOr = std::unexpected(std::current_exception());
    }
 
    std::optional<std::expected<T, std::exception_ptr>> ErrorOr;
};
 
template<typename T>
struct TFutureBase {
    TFutureBase() = default;
    TFutureBase(TPromise<T>& promise)
        : Coro(Coro.from_promise(promise))
    { }
    TFutureBase(TFutureBase&& other)
    {
        *this = std::move(other);
    }
    TFutureBase(const TFutureBase&) = delete;
    TFutureBase& operator=(const TFutureBase&) = delete;
    TFutureBase& operator=(TFutureBase&& other) {
        if (this != &other) {
            if (Coro) {
                Coro.destroy();
            }
            Coro = std::exchange(other.Coro, nullptr);
        }
        return *this;
    }
 
    ~TFutureBase() { if (Coro) { Coro.destroy(); } }
 
    bool await_ready() const {
        return Coro.promise().ErrorOr.has_value();
    }
 
    bool done() const {
        return Coro.done();
    }
 
    std::coroutine_handle<> raw() {
        return Coro;
    }
 
    void await_suspend(std::coroutine_handle<> caller) {
        Coro.promise().Caller = caller;
    }
 
    using promise_type = TPromise<T>;
 
protected:
    std::coroutine_handle<TPromise<T>> Coro = nullptr;
};
 
template<> struct TFuture<void>;
 
template<typename T>
struct TFuture : public TFutureBase<T> {
    T await_resume() {
        auto& errorOr = *this->Coro.promise().ErrorOr;
        if (errorOr.has_value()) {
            return std::move(errorOr.value());
        } else {
            std::rethrow_exception(errorOr.error());
        }
    }
 
 
    template<typename Func>
    auto Apply(Func func) -> TFuture<decltype(func(std::declval<T>()))> {
        auto prev = std::move(*this);
        auto ret = co_await prev;
        co_return func(ret);
    }
 
    TFuture<void> Ignore();
};
 
template<>
struct TPromise<void>: public TPromiseBase<void> {
    TFuture<void> get_return_object();
 
    void return_void() {
        ErrorOr = nullptr;
    }
 
    void unhandled_exception() {
        ErrorOr = std::current_exception();
    }
 
    std::optional<std::exception_ptr> ErrorOr;
};
 
template<>
struct TFuture<void> : public TFutureBase<void> {
    void await_resume() {
        auto& errorOr = *this->Coro.promise().ErrorOr;
        if (errorOr) {
            std::rethrow_exception(errorOr);
        }
    }
 
    template<typename Func>
    auto Accept(Func func) -> TFuture<void> {
        auto prev = std::move(*this);
        co_await prev;
        co_return func();
    }
};
 
template<typename T>
inline TFuture<void> TFuture<T>::Ignore() {
    auto prev = std::move(*this);
    co_await prev;
    co_return;
}
 
template<typename T>
struct TFinalAwaiter {
    bool await_ready() noexcept { return false; }
    std::coroutine_handle<> await_suspend(std::coroutine_handle<TPromise<T>> h) noexcept {
        return h.promise().Caller;
    }
    void await_resume() noexcept { }
};
 
inline TFuture<void> TPromise<void>::get_return_object() { return { TFuture<void>{*this} }; }
template<typename T>
TFuture<T> TPromise<T>::get_return_object() { return { TFuture<T>{*this} }; }
 
 
template<typename T>
TFinalAwaiter<T> TPromiseBase<T>::final_suspend() noexcept { return {}; }
 
template<typename T>
TFuture<std::vector<T>> All(std::vector<TFuture<T>>&& futures) {
    auto waiting = std::move(futures);
    std::vector<T> ret; ret.reserve(waiting.size());
    for (auto& f : waiting) {
        ret.emplace_back(std::move(co_await f));
    }
    co_return ret;
}
 
inline TFuture<void> All(std::vector<TFuture<void>>&& futures) {
    auto waiting = std::move(futures);
    for (auto& f : waiting) {
        co_await f;
    }
    co_return;
}
 
template<typename T>
TFuture<T> Any(std::vector<TFuture<T>>&& futures) {
    std::vector<TFuture<T>> all = std::move(futures);
 
    auto it = std::find_if(all.begin(), all.end(), [](auto& f) { return f.done(); });
    if (it != all.end()) {
        co_return it->await_resume();
    }
 
    auto self = co_await Self();
    for (auto& f : all) {
        f.await_suspend(self);
    }
    co_await std::suspend_always();
    co_return std::find_if(all.begin(), all.end(), [](auto& f) { return f.done(); })->await_resume();
}
 
inline TFuture<void> Any(std::vector<TFuture<void>>&& futures) {
    std::vector<TFuture<void>> all = std::move(futures);
 
    if (std::find_if(all.begin(), all.end(), [](auto& f) { return f.done(); }) != all.end()) {
        co_return;
    }
 
    auto self = co_await Self();
    for (auto& f : all) {
        f.await_suspend(self);
    }
    co_await std::suspend_always();
    co_return;
}
 
} // namespace NNet