generator.hpp

// Copyright Take Vos 2020-2021.
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at https://www.boost.org/LICENSE_1_0.txt)
 
#pragma once
 
#include <concepts>
#include <coroutine>
#include <optional>
#include <memory>
#include <memory_resource>
#include "arguments.hpp"
#include "assert.hpp"
 
namespace hi::inline v1 {
 
template<typename T>
class generator {
public:
    using value_type = T;
 
    class promise_type {
    public:
        generator get_return_object()
        {
            return generator{handle_type::from_promise(*this)};
        }
 
        value_type const &value()
        {
            return *_value;
        }
 
        static std::suspend_never initial_suspend() noexcept
        {
            return {};
        }
 
        static std::suspend_always final_suspend() noexcept
        {
            return {};
        }
 
        std::suspend_always yield_value(value_type const &value) noexcept
        {
            _value = value;
            return {};
        }
 
        std::suspend_always yield_value(value_type &&value) noexcept
        {
            _value = std::move(value);
            return {};
        }
 
        void return_void() noexcept {}
 
        // Disallow co_await in generator coroutines.
        void await_transform() = delete;
 
        void unhandled_exception() noexcept
        {
            _exception = std::current_exception();
        }
 
        void rethrow()
        {
            if (auto ptr = std::exchange(_exception, nullptr)) {
                std::rethrow_exception(ptr);
            }
        }
 
    private:
        std::optional<value_type> _value = {};
        std::exception_ptr _exception = nullptr;
    };
 
    using handle_type = std::coroutine_handle<promise_type>;
 
    class value_proxy {
    public:
        value_proxy(value_type const &value) noexcept : _value(value) {}
 
        value_type const &operator*() const noexcept
        {
            return _value;
        }
 
    private:
        value_type _value;
    };
 
    class iterator {
    public:
        using difference_type = ptrdiff_t;
        using value_type = std::remove_cv_t<value_type>;
        using pointer = value_type *;
        using reference = value_type &;
        using iterator_category = std::input_iterator_tag;
 
        explicit iterator(handle_type coroutine) : _coroutine{coroutine} {}
 
        iterator &operator++()
        {
            hi_axiom(not at_end());
            _coroutine.resume();
            _coroutine.promise().rethrow();
            return *this;
        }
 
        value_proxy operator++(int)
        {
            auto tmp = value_proxy(**this);
            hi_axiom(not at_end());
            _coroutine.resume();
            _coroutine.promise().rethrow();
            return tmp;
        }
 
        value_type const &operator*() const
        {
            hi_axiom(not at_end());
            return _coroutine.promise().value();
        }
 
        value_type const *operator->() const noexcept
        {
            hi_axiom(not at_end());
            return std::addressof(_coroutine.promise().value());
        }
 
        [[nodiscard]] bool at_end() const noexcept
        {
            return (not _coroutine) or _coroutine.done();
        }
 
        [[nodiscard]] bool operator==(std::default_sentinel_t) const noexcept
        {
            return at_end();
        }
 
    private:
        handle_type _coroutine;
    };
 
    explicit generator(handle_type coroutine) : _coroutine(coroutine) {}
 
    generator() = default;
    ~generator()
    {
        if (_coroutine) {
            _coroutine.destroy();
        }
    }
 
    generator(const generator &) = delete;
    generator &operator=(const generator &) = delete;
 
    generator(generator &&other) noexcept : _coroutine{std::exchange(other._coroutine, {})}
    {
    }
 
    generator &operator=(generator &&other) noexcept
    {
        hi_return_on_self_assignment(other);
        if (_coroutine) {
            _coroutine.destroy();
        }
        _coroutine = std::exchange(other._coroutine, {});
        return *this;
    }
 
    iterator begin()
    {
        return iterator{_coroutine};
    }
 
    std::default_sentinel_t end()
    {
        return {};
    }
 
private:
    handle_type _coroutine;
};
 
} // namespace hi::inline v1