generator.hpp

浏览该文件的文档.

#ifndef NEFORCE_CORE_ASYNC_GENERATOR_HPP__
#define NEFORCE_CORE_ASYNC_GENERATOR_HPP__

 
#include "NeForce/core/async/coroutine.hpp"
#if defined(NEFORCE_STANDARD_20) || defined(NEXUSFORCE_ENABLE_DOXYGEN)
#    include "NeForce/core/async/atomic.hpp"
#    include "NeForce/core/exception/exception_ptr.hpp"
#    include "NeForce/core/utility/optional.hpp"
NEFORCE_BEGIN_NAMESPACE__


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

        bool await_ready() const noexcept { return token && token->is_cancelled(); }

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

        void await_resume() const {
            if (token && 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_ && 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_) {
            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_) {
            state_->ref_count.fetch_add(1, memory_order_relaxed);
        }
        return *this;
    }

    ~cancellation_token() { release(); }

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

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

    check_awaiter check() const { return check_awaiter{this}; }
};
 

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 = _NEFORCEmove(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>
    invoke_result_t<F, T> map(F&& func) {
        for (auto&& value: *this) {
            co_yield func(_NEFORCE forward<decltype(value)>(value));
        }
    }

    template <typename Pred>
    generator filter(Pred&& pred) {
        for (auto&& value: *this) {
            if (pred(value)) {
                co_yield _NEFORCE forward<decltype(value)>(value);
            }
        }
    }

    generator take(const size_t n) {
        size_t count = 0;
        for (auto&& value: *this) {
            if (count >= n) {
                break;
            }
            co_yield _NEFORCE forward<decltype(value)>(value);
            ++count;
        }
    }

    generator skip(const size_t n) {
        size_t count = 0;
        for (auto&& value: *this) {
            if (count < n) {
                ++count;
                continue;
            }
            co_yield _NEFORCE forward<decltype(value)>(value);
        }
    }

    generator chain(generator&& other) {
        for (auto&& value: *this) {
            co_yield _NEFORCE forward<decltype(value)>(value);
        }
        for (auto&& value: other) {
            co_yield _NEFORCE forward<decltype(value)>(value);
        }
    }

    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;
    }
};
 

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

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

            void await_resume() noexcept {}
        };
 
        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; }

        bool is_cancelled() const noexcept { return token && 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; 

        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_}; }

    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);
        }
    }

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

template <>
class task<void> {
public:
    struct promise_type {
        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; }
 
        bool is_cancelled() const { return token && token->is_cancelled(); }
 
        suspend_always initial_suspend() noexcept { return {}; }
 
        struct final_awaiter {
            bool await_ready() noexcept { return false; }
 
            coroutine_handle<> await_suspend(coroutine_handle<promise_type> h) noexcept {
                auto& promise = h.promise();
                if (promise.continuation) {
                    return promise.continuation;
                }
                return noop_coroutine();
            }
 
            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;
 
        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_}; }
 
    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);
        }
    }
 
    bool is_cancelled() const { return handle_ && handle_.promise().is_cancelled(); }
};
 

NEFORCE_BEGIN_INNER__
 
template <typename... Ts>
struct when_all_result {
    tuple<Ts...> values;
};
 
template <size_t I, typename Tuple, typename... Tasks, enable_if_t<(I < sizeof...(Tasks)), int> = 0>
task<void> when_all_helper(Tuple& results, Tasks&&... tasks) {
    auto& current_task = _NEFORCE get<I>(_NEFORCE forward_as_tuple(_NEFORCE forward<Tasks>(tasks)...));
    _NEFORCE get<I>(results) = co_await _NEFORCE move(current_task);
    co_await when_all_helper<I + 1>(results, _NEFORCE forward<Tasks>(tasks)...);
    co_return;
}
 
template <size_t I, typename Tuple, typename... Tasks, enable_if_t<I == sizeof...(Tasks), int> = 0>
task<void> when_all_helper(Tuple&, Tasks&&...) {
    co_return;
}
 
NEFORCE_END_INNER__

template <typename... Tasks>
auto when_all(Tasks&&... tasks) -> task<tuple<typename Tasks::promise_type::result_type...>> {
    tuple<typename Tasks::promise_type::result_type...> results;
    co_await inner::when_all_helper<0>(results, _NEFORCE forward<Tasks>(tasks)...);
    co_return results;
}
 

template <typename T, typename Factory>
task<T> retry(Factory&& factory, const size_t max_attempts,
              function<bool(const exception_ptr&)> should_retry = nullptr) {
    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__