generator.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include <coroutine>
#include <exception>
#include <iterator>
#include <utility>
#include <type_traits>
#include <memory>
 
namespace kcenon::container::async
{
    template<typename T>
    class generator;
 
    namespace detail
    {
        template<typename T>
        struct generator_promise
        {
            using value_type = std::remove_reference_t<T>;
            using reference_type = std::conditional_t<std::is_reference_v<T>, T, T&>;
            using pointer_type = value_type*;
 
            pointer_type current_value_{nullptr};
            std::exception_ptr exception_;
 
            [[nodiscard]] generator<T> get_return_object() noexcept;
 
            [[nodiscard]] std::suspend_always initial_suspend() noexcept
            {
                return {};
            }
 
            [[nodiscard]] std::suspend_always final_suspend() noexcept
            {
                return {};
            }
 
            std::suspend_always yield_value(value_type& value) noexcept
            {
                current_value_ = std::addressof(value);
                return {};
            }
 
            std::suspend_always yield_value(value_type&& value) noexcept
            {
                current_value_ = std::addressof(value);
                return {};
            }
 
            void return_void() noexcept {}
 
            void unhandled_exception() noexcept
            {
                exception_ = std::current_exception();
            }
 
            [[nodiscard]] reference_type value() const noexcept
            {
                return *current_value_;
            }
 
            void rethrow_if_exception()
            {
                if (exception_)
                {
                    std::rethrow_exception(exception_);
                }
            }
        };
 
    } // namespace detail
 
    template<typename T>
    class generator
    {
    public:
        using promise_type = detail::generator_promise<T>;
        using handle_type = std::coroutine_handle<promise_type>;
        using value_type = std::remove_reference_t<T>;
        using reference = std::conditional_t<std::is_reference_v<T>, T, T&>;
        using pointer = value_type*;
 
        class iterator
        {
        public:
            using iterator_category = std::input_iterator_tag;
            using difference_type = std::ptrdiff_t;
            using value_type = generator::value_type;
            using reference = generator::reference;
            using pointer = generator::pointer;
 
            iterator() noexcept : handle_(nullptr) {}
 
            explicit iterator(handle_type handle) noexcept : handle_(handle) {}
 
            [[nodiscard]] reference operator*() const noexcept
            {
                return handle_.promise().value();
            }
 
            [[nodiscard]] pointer operator->() const noexcept
            {
                return std::addressof(handle_.promise().value());
            }
 
            iterator& operator++()
            {
                handle_.resume();
                if (handle_.done())
                {
                    handle_.promise().rethrow_if_exception();
                }
                return *this;
            }
 
            void operator++(int)
            {
                ++(*this);
            }
 
            [[nodiscard]] bool operator==(std::default_sentinel_t) const noexcept
            {
                return handle_.done();
            }
 
            [[nodiscard]] bool operator!=(std::default_sentinel_t) const noexcept
            {
                return !handle_.done();
            }
 
        private:
            handle_type handle_;
        };
 
        generator() noexcept : handle_(nullptr) {}
 
        explicit generator(handle_type handle) noexcept : handle_(handle) {}
 
        generator(generator&& other) noexcept
            : handle_(std::exchange(other.handle_, nullptr)) {}
 
        generator& operator=(generator&& other) noexcept
        {
            if (this != &other)
            {
                if (handle_)
                {
                    handle_.destroy();
                }
                handle_ = std::exchange(other.handle_, nullptr);
            }
            return *this;
        }
 
        ~generator()
        {
            if (handle_)
            {
                handle_.destroy();
            }
        }
 
        // Non-copyable
        generator(const generator&) = delete;
        generator& operator=(const generator&) = delete;
 
        [[nodiscard]] bool valid() const noexcept
        {
            return handle_ != nullptr;
        }
 
        [[nodiscard]] explicit operator bool() const noexcept
        {
            return valid();
        }
 
        [[nodiscard]] iterator begin()
        {
            if (handle_)
            {
                handle_.resume();
                if (handle_.done())
                {
                    handle_.promise().rethrow_if_exception();
                }
            }
            return iterator{handle_};
        }
 
        [[nodiscard]] std::default_sentinel_t end() noexcept
        {
            return {};
        }
 
    private:
        handle_type handle_;
    };
 
    // Implementation of get_return_object
    template<typename T>
    generator<T> detail::generator_promise<T>::get_return_object() noexcept
    {
        return generator<T>{
            std::coroutine_handle<detail::generator_promise<T>>::from_promise(*this)};
    }
 
    template<typename Range>
    generator<typename std::ranges::range_value_t<Range>> from_range(Range&& range)
    {
        for (auto&& element : range)
        {
            co_yield std::forward<decltype(element)>(element);
        }
    }
 
    template<typename T>
    generator<T> take(generator<T> gen, size_t count)
    {
        size_t yielded = 0;
        for (auto&& element : gen)
        {
            if (yielded++ >= count)
            {
                break;
            }
            co_yield std::forward<decltype(element)>(element);
        }
    }
 
} // namespace kcenon::container::async