async_generator.h

Go to the documentation of this file.

#ifndef C16AFA99_7780_4436_8BC0_60F24EF771A4
#define C16AFA99_7780_4436_8BC0_60F24EF771A4

 
#include <coroutine>
#include <exception>
#include <functional>
#include <iterator>
#include <memory>
#include <type_traits>
#include <utility>
 
#include "detail.h"
#include "exceptions.h"
#include "task.h"
 
namespace wwa::coro {

template<typename Result>
class [[nodiscard]] async_generator {
public:
    class promise_type {
    public:
        using value_type = std::remove_reference_t<Result>;
        using pointer    = std::add_pointer_t<value_type>;

        promise_type() noexcept = default;

        ~promise_type() = default;

        promise_type(const promise_type&)            = delete;
        promise_type(promise_type&&) noexcept        = delete;
        promise_type& operator=(const promise_type&) = delete;
        promise_type& operator=(promise_type&&)      = delete;

        auto get_return_object() noexcept { return async_generator<Result>{*this}; }

        [[nodiscard]] constexpr auto initial_suspend() const noexcept { return std::suspend_always{}; }

        auto final_suspend() noexcept
        {
            this->m_current_value = nullptr;
            return yield_op{this->m_consumer};
        }

        void unhandled_exception() noexcept { this->m_exception = std::current_exception(); }

        constexpr void return_void() const noexcept {}

        auto yield_value(value_type& value) noexcept
        {
            this->m_current_value = std::addressof(value);
            return yield_op{this->m_consumer};
        }

        auto yield_value(std::add_rvalue_reference_t<std::remove_cv_t<value_type>> value) noexcept
        {
            this->m_current_value = std::addressof(value);
            return yield_op{this->m_consumer};
        }

        [[nodiscard]] constexpr std::add_lvalue_reference_t<value_type> value() const noexcept
        {
            return *this->m_current_value;
        }

        [[nodiscard]] constexpr bool finished() const noexcept { return this->m_current_value == nullptr; }

        void rethrow_if_unhandled_exception()
        {
            if (this->m_exception) {
                std::rethrow_exception(std::move(this->m_exception));
            }
        }

        void set_consumer(std::coroutine_handle<> consumer) noexcept { this->m_consumer = consumer; }
 
    private:
        pointer m_current_value = nullptr;
        std::coroutine_handle<> m_consumer;
        std::exception_ptr m_exception = nullptr;
 
        // NOLINTBEGIN(readability-convert-member-functions-to-static)
        class yield_op {
        public:
            constexpr explicit yield_op(std::coroutine_handle<> consumer) noexcept : m_consumer(consumer) {}

            [[nodiscard]] constexpr bool await_ready() const noexcept { return false; }

            [[nodiscard]] constexpr auto await_suspend([[maybe_unused]] std::coroutine_handle<> h) const noexcept
            {
                return this->m_consumer;
            }

            constexpr void await_resume() const noexcept {}
 
        private:
            std::coroutine_handle<> m_consumer;
        };
        // NOLINTEND(readability-convert-member-functions-to-static)

    };

    class iterator {
    public:

        class advance_op {
        public:
            constexpr advance_op() noexcept : m_promise(nullptr), m_producer(nullptr) {}

            explicit advance_op(std::coroutine_handle<promise_type> producer) noexcept
                : m_promise(std::addressof(producer.promise())), m_producer(producer)
            {}

            [[nodiscard]] constexpr bool await_ready() const noexcept { return this->m_promise == nullptr; }

            constexpr std::coroutine_handle<> await_suspend(std::coroutine_handle<> consumer) noexcept
            {
                this->m_promise->set_consumer(consumer);
                return this->m_producer;
            }

            iterator await_resume()
            {
                if (this->m_promise == nullptr) {
                    return iterator{nullptr};
                }
 
                if (this->m_promise->finished()) {
                    this->m_promise->rethrow_if_unhandled_exception();
                    return iterator{nullptr};
                }
 
                return iterator{std::coroutine_handle<promise_type>::from_promise(*this->m_promise)};
            }
 
        private:
            promise_type* m_promise;
            std::coroutine_handle<> m_producer;
        };


        using iterator_concept = std::input_iterator_tag;

        using difference_type = std::ptrdiff_t;


        using value_type = std::remove_reference_t<Result>;
        using reference  = std::add_lvalue_reference_t<value_type>;
        using pointer    = std::add_pointer_t<value_type>;

        explicit iterator(std::coroutine_handle<promise_type> coroutine) noexcept : m_coroutine(coroutine) {}

        [[nodiscard]] advance_op operator++()
        {
            if (detail::is_good_handle(this->m_coroutine)) {
                return advance_op{this->m_coroutine};
            }

            throw bad_result_access{"Incrementing past the end of the generator"};
        }

        reference operator*() const
        {
            if (detail::is_good_handle(this->m_coroutine)) [[likely]] {
                return this->m_coroutine.promise().value();
            }

            throw bad_result_access{"Dereferencing the end of the generator"};
        }

        friend bool operator==(const iterator& it, [[maybe_unused]] std::default_sentinel_t sentinel) noexcept
        {
            return !detail::is_good_handle(it.m_coroutine);
        }

        bool operator==(const iterator& other) const noexcept
        {
            return this->m_coroutine == other.m_coroutine ||
                   (!detail::is_good_handle(this->m_coroutine) && !detail::is_good_handle(other.m_coroutine));
        }
 
    private:
        std::coroutine_handle<promise_type> m_coroutine;  
    };

    constexpr async_generator() noexcept = default;

    constexpr async_generator(async_generator&& other) noexcept : m_coroutine(std::exchange(other.m_coroutine, nullptr))
    {}

    ~async_generator()
    {
        if (this->m_coroutine) {
            this->m_coroutine.destroy();
        }
    }

    async_generator& operator=(async_generator&& other) noexcept
    {
        if (this != std::addressof(other)) {
            if (this->m_coroutine) {
                this->m_coroutine.destroy();
            }
 
            this->m_coroutine = std::exchange(other.m_coroutine, nullptr);
        }
 
        return *this;
    }

    async_generator(const async_generator&)            = delete;
    async_generator& operator=(const async_generator&) = delete;

    [[nodiscard]] auto begin() noexcept
    {
        if (!detail::is_good_handle(this->m_coroutine)) {
            return typename iterator::advance_op();
        }
 
        return typename iterator::advance_op{this->m_coroutine};
    }

    [[nodiscard]] constexpr auto end() const noexcept { return std::default_sentinel; }
 
private:
    std::coroutine_handle<promise_type> m_coroutine;  

    explicit async_generator(promise_type& promise) noexcept
        : m_coroutine(std::coroutine_handle<promise_type>::from_promise(promise))
    {}
};
 
template<typename AsyncGenerator, typename Callable>
task<void> async_for_each(AsyncGenerator gen, Callable callable)
{
    auto it  = co_await gen.begin();
    auto end = gen.end();
 
    while (it != end) {
        std::invoke(std::forward<Callable>(callable), *it);
        co_await ++it;
    }
}

 
}  // namespace wwa::coro
 
#endif /* C16AFA99_7780_4436_8BC0_60F24EF771A4 */