generator.hpp

浏览该文件的文档.

#ifndef NEFORCE_CORE_ASYNC_GENERATOR_HPP__
#define NEFORCE_CORE_ASYNC_GENERATOR_HPP__

 
#include "NeForce/core/async/coroutine.hpp"
#ifdef NEFORCE_STANDARD_20
#    include "NeForce/core/async/atomic.hpp"
#    include "NeForce/core/exception/exception_ptr.hpp"
#    include "NeForce/core/functional/function.hpp"
#    include "NeForce/core/utility/optional.hpp"
#    include "NeForce/core/utility/tuple.hpp"
NEFORCE_BEGIN_NAMESPACE__


class cancellation_token {
public:
    struct check_awaiter {
        const cancellation_token* token; 

        NEFORCE_NODISCARD bool await_ready() const noexcept { return token != nullptr && token->is_cancelled(); }

        void await_suspend(coroutine_handle<> h) const noexcept {}

        void await_resume() const {
            if (token != nullptr && token->is_cancelled()) {
                NEFORCE_THROW_EXCEPTION(exception("Operation cancelled"));
            }
        }
    };
 
private:
    struct state {
        atomic<bool> cancelled{false}; 
        atomic<size_t> ref_count{1};   
    };
    state* state_; 

    void release() {
        if (state_ != nullptr && state_->ref_count.fetch_sub(1, memory_order_acq_rel) == 1) {
            delete state_;
        }
    }
 
public:
    cancellation_token() :
    state_(new state()) {}

    cancellation_token(const cancellation_token& other) :
    state_(other.state_) {
        if (state_ != nullptr) {
            state_->ref_count.fetch_add(1, memory_order_relaxed);
        }
    }

    cancellation_token& operator=(const cancellation_token& other) {
        if (addressof(other) == this) {
            return *this;
        }
 
        release();
        state_ = other.state_;
        if (state_ != nullptr) {
            state_->ref_count.fetch_add(1, memory_order_relaxed);
        }
        return *this;
    }

    ~cancellation_token() { release(); }

    void cancel() noexcept {
        if (state_ != nullptr) {
            state_->cancelled.store(true, memory_order_release);
        }
    }

    NEFORCE_NODISCARD bool is_cancelled() const noexcept {
        return state_ != nullptr && state_->cancelled.load(memory_order_acquire);
    }

    NEFORCE_NODISCARD check_awaiter check() const { return check_awaiter{this}; }
};
 
 
template <typename T>
class generator;

NEFORCE_BEGIN_INNER__
 
template <typename T, typename F>
generator<invoke_result_t<F, T>> generator_map(generator<T> source, F func);
 
template <typename T, typename Pred>
generator<T> generator_filter(generator<T> source, Pred pred);
 
template <typename T>
generator<T> generator_take(generator<T> source, size_t n);
 
template <typename T>
generator<T> generator_skip(generator<T> source, size_t n);
 
template <typename T>
generator<T> generator_chain(generator<T> source, generator<T> other);
 
NEFORCE_END_INNER__
 

template <typename T>
class generator {
public:
    struct promise_type {
        optional<T> current_value; 
        exception_ptr exception;   

        generator get_return_object() { return generator{coroutine_handle<promise_type>::from_promise(*this)}; }

        suspend_always initial_suspend() noexcept { return {}; }

        suspend_always final_suspend() noexcept { return {}; }

        suspend_always yield_value(T value) {
            current_value = _NEFORCE move(value);
            return {};
        }

        void return_void() noexcept {}

        void unhandled_exception() { exception = _NEFORCE current_exception(); }
    };

    struct iterator {
        using iterator_category = input_iterator_tag;
        using value_type = T;
        using difference_type = ptrdiff_t;
        using pointer = T*;
        using reference = T&;
 
        coroutine_handle<promise_type> handle; 

        iterator() = default;

        explicit iterator(coroutine_handle<promise_type> h) :
        handle(h) {}

        iterator& operator++() {
            handle.resume();
            if (handle.done()) {
                handle = nullptr;
            }
            return *this;
        }

        void operator++(int) { ++(*this); }

        T& operator*() const noexcept {
            NEFORCE_CONSTEXPR_ASSERT(handle.promise().current_value.has_value());
            return *handle.promise().current_value;
        }

        T* operator->() const noexcept { return &(*handle.promise().current_value); }

        bool operator==(const iterator& other) const noexcept { return handle == other.handle; }

        bool operator!=(const iterator& other) const noexcept { return !(*this == other); }
    };
 
private:
    coroutine_handle<promise_type> handle_; 
 
public:
    explicit generator(coroutine_handle<promise_type> h) :
    handle_(h) {}

    generator(generator&& other) noexcept :
    handle_(other.handle_) {
        other.handle_ = nullptr;
    }

    generator& operator=(generator&& other) noexcept {
        if (addressof(other) == this) {
            return *this;
        }
 
        if (handle_) {
            handle_.destroy();
        }
        handle_ = other.handle_;
        other.handle_ = nullptr;
 
        return *this;
    }

    ~generator() {
        if (handle_) {
            handle_.destroy();
        }
    }
 
    generator(const generator&) = delete;
    generator& operator=(const generator&) = delete;

    iterator begin() {
        if (handle_) {
            handle_.resume();
            if (handle_.promise().exception) {
                _NEFORCE rethrow_exception(handle_.promise().exception);
            }
            if (handle_.done()) {
                return iterator{nullptr};
            }
        }
        return iterator{handle_};
    }

    iterator end() { return iterator{nullptr}; }

    template <typename F>
    generator<invoke_result_t<F, T>> map(F func) {
        return inner::generator_map(move(*this), move(func));
    }

    template <typename Pred>
    generator filter(Pred pred) {
        return inner::generator_filter(move(*this), move(pred));
    }

    generator take(const size_t n) { return inner::generator_take(move(*this), n); }

    generator skip(const size_t n) { return inner::generator_skip(move(*this), n); }

    generator chain(generator other) { return inner::generator_chain(move(*this), move(other)); }

    template <typename F>
    void for_each(F&& func) {
        for (auto&& value: *this) {
            func(_NEFORCE forward<decltype(value)>(value));
        }
    }

    template <typename Acc, typename F>
    Acc fold(Acc init, F&& func) {
        Acc result = _NEFORCE move(init);
        for (auto&& value: *this) {
            result = func(_NEFORCE move(result), _NEFORCE forward<decltype(value)>(value));
        }
        return result;
    }
};
 

NEFORCE_BEGIN_INNER__
 
template <typename T, typename F>
generator<invoke_result_t<F, T>> generator_map(generator<T> source, F func) {
    for (auto&& value: source) {
        co_yield func(_NEFORCE forward<decltype(value)>(value));
    }
}
 
template <typename T, typename Pred>
generator<T> generator_filter(generator<T> source, Pred pred) {
    for (auto&& value: source) {
        if (pred(value)) {
            co_yield _NEFORCE forward<decltype(value)>(value);
        }
    }
}
 
template <typename T>
generator<T> generator_take(generator<T> source, size_t n) {
    size_t count = 0;
    for (auto&& value: source) {
        if (count >= n) {
            break;
        }
        co_yield _NEFORCE forward<decltype(value)>(value);
        ++count;
    }
}
 
template <typename T>
generator<T> generator_skip(generator<T> source, size_t n) {
    size_t count = 0;
    for (auto&& value: source) {
        if (count < n) {
            ++count;
            continue;
        }
        co_yield _NEFORCE forward<decltype(value)>(value);
    }
}
 
template <typename T>
generator<T> generator_chain(generator<T> source, generator<T> other) {
    for (auto&& value: source) {
        co_yield _NEFORCE forward<decltype(value)>(value);
    }
    for (auto&& value: other) {
        co_yield _NEFORCE forward<decltype(value)>(value);
    }
}
 
NEFORCE_END_INNER__

template <typename T>
class task {
public:
    struct promise_type {
        struct final_awaiter {
            bool await_ready() noexcept { return false; }

            void await_suspend(coroutine_handle<promise_type> h) noexcept {
                auto& promise = h.promise();
                if (promise.continuation) {
                    promise.continuation.resume();
                }
            }

            void await_resume() noexcept {}
        };
 
        using value_type = T; 
 
        optional<T> result;                  
        exception_ptr exception;             
        coroutine_handle<> continuation;     
        cancellation_token* token = nullptr; 

        task get_return_object() { return task{coroutine_handle<promise_type>::from_promise(*this)}; }

        void set_cancellation_token(cancellation_token* t) noexcept { token = t; }

        NEFORCE_NODISCARD bool is_cancelled() const noexcept { return token != nullptr && token->is_cancelled(); }

        suspend_always initial_suspend() noexcept { return {}; }

        final_awaiter final_suspend() noexcept { return {}; }

        void return_value(T value) { result = _NEFORCE move(value); }

        void unhandled_exception() noexcept { exception = current_exception(); }
    };

    struct awaiter {
        coroutine_handle<promise_type> handle; 

        NEFORCE_NODISCARD bool await_ready() const noexcept { return handle.done(); }

        coroutine_handle<> await_suspend(coroutine_handle<> continuation) noexcept {
            handle.promise().continuation = continuation;
            return handle;
        }

        T await_resume() {
            if (handle.promise().exception) {
                _NEFORCE rethrow_exception(handle.promise().exception);
            }
            return _NEFORCE move(*handle.promise().result);
        }
    };
 
private:
    coroutine_handle<promise_type> handle_; 
 
public:
    explicit task(coroutine_handle<promise_type> h) :
    handle_(h) {}

    task(task&& other) noexcept :
    handle_(other.handle_) {
        other.handle_ = nullptr;
    }

    task& operator=(task&& other) noexcept {
        if (addressof(other) == this) {
            return *this;
        }
 
        if (handle_) {
            handle_.destroy();
        }
        handle_ = other.handle_;
        other.handle_ = nullptr;
 
        return *this;
    }

    ~task() {
        if (handle_) {
            handle_.destroy();
        }
    }
 
    task(const task&) = delete;
    task& operator=(const task&) = delete;

    awaiter operator co_await() { return awaiter{handle_}; }

    NEFORCE_NODISCARD bool done() const noexcept { return handle_.done(); }

    void resume() {
        if (handle_ && !handle_.done()) {
            handle_.resume();
        }
    }

    T get() {
        while (!handle_.done()) {
            handle_.resume();
        }
        if (handle_.promise().exception) {
            _NEFORCE rethrow_exception(handle_.promise().exception);
        }
        return _NEFORCE move(*handle_.promise().result);
    }

    void set_cancellation_token(cancellation_token& token) {
        if (handle_) {
            handle_.promise().set_cancellation_token(&token);
        }
    }

    NEFORCE_NODISCARD bool is_cancelled() const noexcept { return handle_ && handle_.promise().is_cancelled(); }
};
 

template <>
class task<void> {
public:
    struct promise_type {
        using value_type = void; 
 
        exception_ptr exception;
        coroutine_handle<> continuation;
        cancellation_token* token = nullptr;
 
        task get_return_object() { return task{coroutine_handle<promise_type>::from_promise(*this)}; }
 
        void set_cancellation_token(cancellation_token* t) { token = t; }
 
        NEFORCE_NODISCARD bool is_cancelled() const { return token != nullptr && token->is_cancelled(); }
 
        suspend_always initial_suspend() noexcept { return {}; }
 
        struct final_awaiter {
            bool await_ready() noexcept { return false; }
 
            void await_suspend(coroutine_handle<promise_type> h) noexcept {
                auto& promise = h.promise();
                if (promise.continuation) {
                    promise.continuation.resume();
                }
            }
 
            void await_resume() noexcept {}
        };
 
        final_awaiter final_suspend() noexcept { return {}; }
 
        void return_void() noexcept {}
 
        void unhandled_exception() { exception = current_exception(); }
    };
 
    struct awaiter {
        coroutine_handle<promise_type> handle;
 
        NEFORCE_NODISCARD bool await_ready() const noexcept { return handle.done(); }
 
        coroutine_handle<> await_suspend(coroutine_handle<> continuation) noexcept {
            handle.promise().continuation = continuation;
            return handle;
        }
 
        void await_resume() {
            if (handle.promise().exception) {
                rethrow_exception(handle.promise().exception);
            }
        }
    };
 
private:
    coroutine_handle<promise_type> handle_;
 
public:
    explicit task(coroutine_handle<promise_type> h) :
    handle_(h) {}
 
    task(task&& other) noexcept :
    handle_(other.handle_) {
        other.handle_ = nullptr;
    }
 
    task& operator=(task&& other) noexcept {
        if (addressof(other) == this) {
            return *this;
        }
 
        if (handle_) {
            handle_.destroy();
        }
        handle_ = other.handle_;
        other.handle_ = nullptr;
 
        return *this;
    }
 
    ~task() {
        if (handle_) {
            handle_.destroy();
        }
    }
 
    task(const task&) = delete;
    task& operator=(const task&) = delete;
 
    awaiter operator co_await() { return awaiter{handle_}; }
 
    NEFORCE_NODISCARD bool done() const { return handle_.done(); }
 
    void resume() {
        if (handle_ && !handle_.done()) {
            handle_.resume();
        }
    }
 
    void get() {
        while (!handle_.done()) {
            handle_.resume();
        }
        if (handle_.promise().exception) {
            rethrow_exception(handle_.promise().exception);
        }
    }
 
    void set_cancellation_token(cancellation_token& token) {
        if (handle_) {
            handle_.promise().set_cancellation_token(&token);
        }
    }
 
    NEFORCE_NODISCARD bool is_cancelled() const { return handle_ && handle_.promise().is_cancelled(); }
};
 

template <typename... Tasks>
auto when_all(Tasks... tasks) -> task<tuple<typename Tasks::promise_type::value_type...>> {
    tuple<typename Tasks::promise_type::value_type...> results;
    auto task_tuple = _NEFORCE make_tuple(_NEFORCE move(tasks)...);
    [&]<size_t... Is>(index_sequence<Is...>) {
        ((_NEFORCE get<Is>(results) = _NEFORCE get<Is>(task_tuple).get()), ...);
    }(index_sequence_for<Tasks...>{});
    co_return results;
}
 

template <typename T, typename Factory>
task<T> retry(Factory factory, const size_t max_attempts,
              const function<bool(const exception_ptr&)>& should_retry = {}) {
    static_assert(_NEFORCE is_invocable_r_v<task<T>, Factory>, "Factory must return task<T>");
 
    exception_ptr last_exception;
 
    for (size_t attempt = 0; attempt < max_attempts; ++attempt) {
        try {
            co_return co_await factory();
        } catch (...) {
            last_exception = _NEFORCE current_exception();
 
            if (should_retry && !should_retry(last_exception)) {
                rethrow_exception(last_exception);
            }
            if (attempt == max_attempts - 1) {
                rethrow_exception(last_exception);
            }
        }
    }
 
    rethrow_exception(last_exception);
}
 // Coroutine
 
NEFORCE_END_NAMESPACE__
#endif // NEFORCE_STANDARD_20
#endif // NEFORCE_CORE_ASYNC_GENERATOR_HPP__