generator.hpp

Go to the documentation of this file.

#pragma once
 
#include <cortex/coroutine.hpp>
#include <cortex/coroutine_suspend_context.hpp>
#include <cortex/memory_resource.hpp>
 
#include <function2/function2.hpp>
 
#include <cassert>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>
 
namespace cortex {
 
template <typename T>
class Generator final {
public:
private:
    struct State;
 
public:
    class YieldContext {
    public:
        YieldContext(const YieldContext&) = delete;
        YieldContext(YieldContext&&) = delete;
        YieldContext& operator=(const YieldContext&) = delete;
        YieldContext& operator=(YieldContext&&) = delete;
 
        template <typename U>
        void operator()(U&& value) {
            assert(state_);
            assert(ctx_);
            assert(!state_->current.has_value());
            state_->current.emplace(std::forward<U>(value));
            ctx_->Suspend();
        }
 
    private:
        friend class Generator;
 
        YieldContext(State* state, CoroutineSuspendContext* ctx) noexcept
            : state_(state)
            , ctx_(ctx) {}
 
        State* state_ {nullptr};
        CoroutineSuspendContext* ctx_ {nullptr};
    };
 
    using Body = fu2::unique_function<void(YieldContext&)>;
 
    static Generator Make(Body body,
                          std::size_t stack_size_bytes = 262144,
                          MemoryResourceSharedPtr resource = GetDefaultMemoryResource()) {
        if (!static_cast<bool>(body)) {
            throw std::invalid_argument("generator body is null.");
        }
 
        auto state = std::make_shared<State>();
        auto coroutine = Coroutine::Make(
            [state, body = std::move(body)](CoroutineSuspendContext& ctx) mutable {
                YieldContext yield_ctx(state.get(), &ctx);
                body(yield_ctx);
            },
            stack_size_bytes,
            std::move(resource));
 
        return Generator(std::move(coroutine), std::move(state));
    }
 
    struct Builder {
    public:
        Builder()
            : stack_size_bytes_(262144)
            , memory_resource_(GetDefaultMemoryResource()) {}
 
        Generator Build(Body body) && {
            return Generator::Make(std::move(body), stack_size_bytes_, std::move(memory_resource_));
        }
 
        Builder SetStackSizeInBytes(std::size_t stack_size_bytes) && noexcept {
            stack_size_bytes_ = stack_size_bytes;
            return std::move(*this);
        }
 
        Builder SetMemoryResource(MemoryResourceSharedPtr resource) && noexcept {
            memory_resource_ = std::move(resource);
            return std::move(*this);
        }
 
    private:
        std::size_t stack_size_bytes_ {0};
        MemoryResourceSharedPtr memory_resource_ {nullptr};
    };
 
    Generator(const Generator&) = delete;
    Generator(Generator&&) noexcept = default;
    Generator& operator=(const Generator&) = delete;
    Generator& operator=(Generator&&) noexcept = default;
    ~Generator() = default;
 
    [[nodiscard]] bool IsDone() const noexcept {
        return coroutine_.IsDone();
    }
 
    bool Next() {
        assert(state_);
        if (coroutine_.IsDone()) {
            return false;
        }
 
        state_->current.reset();
 
        coroutine_.Resume();
        return state_->current.has_value();
    }
 
    T DetachValue() {
        assert(state_);
        if (!state_->current.has_value()) {
            throw std::logic_error("generator has no value.");
        }
 
        T value = std::move(*state_->current);
        state_->current.reset();
        return value;
    }
 
private:
    struct State {
        std::optional<T> current;
    };
 
    explicit Generator(Coroutine coroutine, std::shared_ptr<State> state)
        : state_(std::move(state))
        , coroutine_(std::move(coroutine)) {}
 
    std::shared_ptr<State> state_;
    Coroutine coroutine_;
};
 
} // namespace cortex