generator.h

#pragma once
 
#include "allocator.h"
#include "future.h"
 
namespace coro {
 
 
template<typename Src>
class generator_iterator {
public:
 
 
    using storage_type = std::invoke_result_t<Src>;
 
    using iterator_category = std::input_iterator_tag;
    using value_type = typename std::decay_t<Src>::value_type;
    using reference = future<value_type> &;
    using difference_type = std::ptrdiff_t;
    using pointer = std::add_pointer_t<value_type>;
 
    reference operator*() {
        return _stor;
    }
 
 
    generator_iterator &operator++() {
        _stor = _src();
        _stor.start();
        if (!_stor.is_pending() && !_stor.has_value()) {
            _is_end = true;
        }
        return *this;
    }
 
    bool operator==(const generator_iterator &other) const {
        return _is_end == other._is_end;
    }
 
    static generator_iterator begin(Src src) {
        generator_iterator ret(src, false);
        ++ret;
        return ret;
    }
 
    static generator_iterator end(Src src) {
        return generator_iterator(src, true);
    }
 
protected:
    Src _src;
    bool _is_end;
    mutable storage_type _stor;
 
    generator_iterator(Src src, bool is_end):_src(src), _is_end(is_end) {
 
    }
 
};
 
 
template<typename T, coro_allocator Alloc = std_allocator>
class generator {
public:
 
    using value_type = T;
    using reference = std::add_lvalue_reference_t<T>;
 
    class promise_type: public _details::coro_promise_base<T>,
                        public coro_allocator_helper<Alloc> {
    public:
 
        promise_type() {
            trace::set_class(std::coroutine_handle<promise_type>::from_promise(*this), typeid(generator).name());
        }
 
        constexpr trace::suspend_always initial_suspend() const {return {};}
 
        struct switch_awaiter {
            static constexpr bool await_ready() noexcept {return false;}
            static constexpr void await_resume() noexcept {}
 
            std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h, std::source_location loc = std::source_location::current()) noexcept {
                promise_type &self = h.promise();
                return trace::on_switch(h, self.set_resolved_switch(), &loc);
            }
        };
 
        struct final_switch_awaiter {
            static constexpr bool await_ready() noexcept {return false;}
            static constexpr void await_resume() noexcept {}
 
            std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept {
                promise_type &self = h.promise();
                return trace::on_switch(h, self.set_resolved_switch(), {});
            }
        };
 
        /* WORKAROUND: clang 16+ fails here resulting in code crash during test - this is workaround
         *
         * CORO_OPT_BARRIER separates optimization between suspended coroutine and
         * normal code. For the clang it means to disable optimizations
         *
         * DETAILS: return value of set_resolved is prepared_coro which
         * must be initialized in current stack frame. However, the clang's
         * too aggresive optimization causes that this value is initialized
         * in coroutine frame at the same address as next awaiter, which
         * leads to overwritting its content on next (probably nested)
         * suspend event. The code crashes when it returns from suspenssion.
         *
         */
 
        CORO_OPT_BARRIER inline std::coroutine_handle<> set_resolved_switch() {
            return this->set_resolved().symmetric_transfer();
        }
 
 
        template<typename Arg>
        requires future_constructible<T, Arg>
        switch_awaiter yield_value(Arg && val) {
            trace::on_yield(std::coroutine_handle<const promise_type>::from_promise(*this), val);
            this->set_value(std::forward<Arg>(val));
            return {};
        }
        switch_awaiter yield_value(std::exception_ptr e) {
            trace::on_yield(std::coroutine_handle<const promise_type>::from_promise(*this), e);
            this->set_exception(std::move(e));
            return {};
        }
 
        final_switch_awaiter final_suspend() noexcept {
            return {};
        }
 
        void return_void() {
        }
 
        bool done() const {
            return std::coroutine_handle<const promise_type>::from_promise(*this).done();
        }
 
        generator get_return_object() {return this;}
 
        deferred_future<T> resume() {
            return [this](auto promise) -> std::coroutine_handle<> {
                if (done()) return {};
                this->fut = promise.release();
                auto h = std::coroutine_handle<promise_type>::from_promise(*this);
                trace::awaiting_ref(h, this->fut);
                return h;
            };
        }
    };
 
 
    deferred_future<T> operator()() {
        return _prom->resume();
    }
 
 
 
    explicit operator bool() const {
        return !_prom->done();
    }
 
    generator() = default;
 
    template<typename A>
    generator(generator<T, A> &&other): _prom(cast_promise(other._prom.release())) {
 
    }
 
 
    auto begin() {
        return generator_iterator<generator &>::begin(*this);
    }
 
    auto end() {
        return generator_iterator<generator &>::end(*this);
    }
 
    operator ident_t() const {return std::coroutine_handle<promise_type>::from_promise(*_prom);}
 
protected:
 
    template<typename X>
    static promise_type *cast_promise(X *other) {
        return static_cast<promise_type *>(static_cast<_details::coro_promise_base<T> *>(other));
    }
 
    generator(promise_type *p):_prom(p) {}
 
    struct deleter {
        void operator()(promise_type *x) {
            auto h = std::coroutine_handle<promise_type>::from_promise(*x);
            h.destroy();
        }
    };
 
    template<typename A, coro_allocator B> friend class generator;
 
    std::unique_ptr<promise_type,deleter> _prom;
 
};
 
template<typename R, coro_allocator Alloc>
class generator<R(), Alloc>: public generator<R, Alloc> {
public:
    using generator<R, Alloc>::generator;
};
 
class fetch_args_tag {};
 
inline constexpr fetch_args_tag fetch_args = {};
 
 
template<typename R, coro_allocator Alloc, typename ... Args>
class generator<R(Args...), Alloc> {
public:
 
    using return_type = R;
 
    using arg_type = std::tuple<Args...>;
 
    class promise_type: public _details::coro_promise_base<R>,
                        public coro_allocator_helper<Alloc> {
    public:
 
        promise_type() {
            trace::set_class(std::coroutine_handle<promise_type>::from_promise(*this), typeid(generator).name());
        }
 
 
        constexpr trace::suspend_always initial_suspend() const {return {};}
 
        struct fetch_arg_awaiter {
            arg_type *arg;
            static constexpr bool await_ready() noexcept {return true;}
            static constexpr void await_suspend(std::coroutine_handle<>) noexcept {}
            arg_type &await_resume() noexcept {return *arg;};
        };
 
        struct switch_awaiter {
            promise_type *self = nullptr;
            static constexpr bool await_ready() noexcept {return false;}
            arg_type &await_resume() noexcept {return *self->arg;}
 
            CORO_OPT_BARRIER std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h, std::source_location loc = std::source_location::current()) noexcept {
                self = &h.promise();
                return trace::on_switch(h,self->set_resolved().symmetric_transfer(),&loc);
            }
        };
 
        struct switch_final_awaiter: public switch_awaiter {
            CORO_OPT_BARRIER std::coroutine_handle<> await_suspend(std::coroutine_handle<promise_type> h) noexcept {
                this->self = &h.promise();
                return trace::on_switch(h,this->self->set_resolved().symmetric_transfer(),{});
            }
        };
 
 
        template<typename Arg>
        requires future_constructible<R, Arg>
        switch_awaiter yield_value(Arg && val) {
            trace::on_yield(std::coroutine_handle<const promise_type>::from_promise(*this), val);
            this->set_value(std::forward<Arg>(val));
            return {};
        }
 
        switch_awaiter yield_value(std::exception_ptr e) {
            LIBCORO_TRACE_YIELD(std::coroutine_handle<const promise_type>::from_promise(*this),e);
            this->set_exception(std::move(e));
            return {};
        }
 
 
        fetch_arg_awaiter yield_value(std::nullptr_t) {
            return {&(*arg)};
        }
 
        fetch_arg_awaiter yield_value(fetch_args_tag) {
            return {&(*arg)};
        }
 
        switch_final_awaiter final_suspend() noexcept {
            return {};
        }
 
        void return_void() {
        }
 
        bool done() const {
            return std::coroutine_handle<const promise_type>::from_promise(*this).done();
        }
 
        generator get_return_object() {return this;}
 
        future<R> resume() {
            return [this](auto promise)  {
                if (done()) return;
                this->fut = promise.release();
                auto h = std::coroutine_handle<promise_type>::from_promise(*this);
                trace::awaiting_ref(h, this->fut);
                trace::resume(h);
            };
        }
 
        std::optional<arg_type> arg;
 
        template<typename ... XArgs>
        void set_arg(XArgs && ... args) {
            arg.emplace(std::forward<XArgs>(args)...);
        }
    };
 
    explicit operator bool() const {
        return !_prom->done();
    }
 
    generator() = default;
 
    template<typename A>
    generator(generator<R(Args...), A> &&other): _prom(cast_promise(other._prom.release())) {
 
    }
 
 
    template<typename ... XArgs>
    future<R> operator()(XArgs &&... args) {
        static_assert(std::is_constructible_v<arg_type, XArgs...>);
        _prom->set_arg(std::forward<XArgs>(args)...);
        return _prom->resume();
    }
 
 
 
protected:
 
    template<typename X>
    static promise_type *cast_promise(X *other) {
        return static_cast<promise_type *>(static_cast<_details::coro_promise_base<R> *>(other));
    }
 
    generator(promise_type *p):_prom(p) {}
 
    struct deleter {
        void operator()(promise_type *x) {
            auto h = std::coroutine_handle<promise_type>::from_promise(*x);
            h.destroy();
        }
    };
 
    template<typename A, coro_allocator B> friend class generator;
 
    std::unique_ptr<promise_type,deleter> _prom;
 
 
 
};
 
 
 
}