diff options
Diffstat (limited to 'src/libutil/generator.hh')
| -rw-r--r-- | src/libutil/generator.hh | 289 |
1 files changed, 289 insertions, 0 deletions
diff --git a/src/libutil/generator.hh b/src/libutil/generator.hh new file mode 100644 index 000000000..fda257002 --- /dev/null +++ b/src/libutil/generator.hh @@ -0,0 +1,289 @@ +#pragma once +///@file + +#include "types.hh" + +#include <coroutine> +#include <exception> +#include <optional> +#include <utility> +#include <variant> + +namespace nix { + +template<typename T, typename Transform> +struct Generator; + +namespace _generator { + +template<typename T> +struct promise_state; +template<typename T> +struct GeneratorBase; + +struct finished {}; + +template<typename T> +struct link +{ + std::coroutine_handle<> handle{}; + promise_state<T> * state{}; +}; + +struct failure +{ + std::exception_ptr e; +}; + +template<typename T> +struct promise_state +{ + // result of the most recent coroutine resumption: a value, + // a nested coroutine to drain, an error, or our completion + std::variant<T, link<T>, failure, finished> value{}; + // coroutine to resume when this one has finished. set when + // one generator yields another, such that the entire chain + // of parents always linearly points to the root generator. + link<T> parent{}; +}; + +template<typename T, typename Transform> +struct promise : promise_state<T> +{ + using transform_t = std::conditional_t<std::is_void_v<Transform>, std::identity, Transform>; + + transform_t convert; + std::optional<GeneratorBase<T>> inner; + + // called by the compiler to convert the internal promise object + // to the user-declared (function return) type of the coroutine. + Generator<T, Transform> get_return_object() + { + auto h = std::coroutine_handle<promise>::from_promise(*this); + return Generator<T, Transform>(GeneratorBase<T>(h, h.promise())); + } + std::suspend_always initial_suspend() + { + return {}; + } + std::suspend_always final_suspend() noexcept + { + return {}; + } + void unhandled_exception() + { + this->value = failure{std::current_exception()}; + } + + // `co_yield` handler for "simple" values, i.e. those that + // are transformed directly to a T by the given transform. + template<typename From> + requires requires(transform_t t, From && f) { + { + t(std::forward<From>(f)) + } -> std::convertible_to<T>; + } + std::suspend_always yield_value(From && from) + { + this->value.template emplace<0>(convert(std::forward<From>(from))); + return {}; + } + + // `co_yield` handler for "complex" values, i.e. those that + // are transformed into another generator. we'll drain that + // new generator completely before resuming the current one + template<typename From> + requires requires(transform_t t, From && f) { + static_cast<Generator<T, void>>(t(std::forward<From>(f))); + } + std::suspend_always yield_value(From && from) + { + inner = static_cast<Generator<T, void>>(convert(std::forward<From>(from))).impl; + this->value = inner->active; + return {}; + } + + // handler for `co_return`, including the implicit `co_return` + // at the end of a coroutine that does not have one explicitly + void return_void() + { + this->value = finished{}; + } +}; + +template<typename T> +struct GeneratorBase +{ + template<typename, typename> + friend struct Generator; + template<typename, typename> + friend struct promise; + + // NOTE coroutine handles are LiteralType, own a memory resource (that may + // itself own unique resources), and are "typically TriviallyCopyable". we + // need to take special care to wrap this into a less footgunny interface. + GeneratorBase(GeneratorBase && other) + { + swap(other); + } + + GeneratorBase & operator=(GeneratorBase && other) + { + GeneratorBase(std::move(other)).swap(*this); + return *this; + } + + ~GeneratorBase() + { + if (h) { + h.destroy(); + } + } + + std::optional<T> next() + { + // resume the currently active coroutine once. it can return either a + // value, an exception, another generator to drain, or it can finish. + // since c++ coroutines cannot directly return anything from resume() + // we must communicate all results via `active->state.value` instead. + while (active.handle) { + active.handle.resume(); + auto & p = *active.state; + // process the result. only one case sets this to a non-`nullopt` + // value, all others leave it at `nullopt` to request more loops. + auto result = std::visit( + overloaded{ + // when the current coroutine handle is done we'll try to + // resume its parent (if the current handle was retrieved + // from a `co_yield`ed generator) or finish the generator + // entirely because the root active.parent has no handle. + [&](finished) -> std::optional<T> { + active = p.parent; + return {}; + }, + // when the coroutine yields a generator we push the full + // inner stack onto our own stack and resume the top item + [&](link<T> & inner) -> std::optional<T> { + auto base = inner.state; + while (base->parent.handle) { + base = base->parent.state; + } + base->parent = active; + active = inner; + return {}; + }, + // values are simply returned to the caller, as received. + [&](T & value) -> std::optional<T> { return std::move(value); }, + // exceptions must be rethrown. resuming again after this + // is not allowed because the top-most coroutine would be + // finished and we'd thus step back to its parent, but by + // doing so we might invalidate invariants of the parent. + // allowing the parent to catch exceptions of a child for + // `co_yield` exceptions specifically would introduce far + // too many problems to be worth the doing (since parents + // can neither know nor revert any yields of their child) + [&](failure & f) -> std::optional<T> { + active = {}; + std::rethrow_exception(f.e); + }, + }, + p.value + ); + if (result) { + return result; + } + } + + return std::nullopt; + } + +protected: + std::coroutine_handle<> h{}; + link<T> active{}; + + GeneratorBase(std::coroutine_handle<> h, promise_state<T> & state) + : h(h) + , active(h, &state) + { + } + + void swap(GeneratorBase & other) + { + std::swap(h, other.h); + std::swap(active, other.active); + } +}; + +} // _generator + +/// Coroutine-based iterator modeled loosely on Rust [`std::iter::Iterator`][iter] +/// interface. Like Rust's `Iterator` and unlike common C++ iterators, a Generator +/// returns `std::optional<T>` values from its next() function, but unlike both it +/// can also transform items produced within using a Transform function object the +/// Generator holds before returning them via next(). To allow generator nesting a +/// Transform may also return another Generator instance for any yielded value, in +/// this case the new Generator will temporarily take priority over the previously +/// running one and have its values returned until it is exhausted, then return to +/// the previous Generator. This mechanism may nest Generator to arbitrary depths. +/// +/// \tparam T item type +/// \tparam Transform transform function object type, or `void` for no transform +/// +/// [iter]: https://doc.rust-lang.org/stable/std/iter/trait.Iterator.html +template<typename T, typename Transform = void> +struct Generator +{ + template<typename, typename> + friend struct _generator::promise; + // erasing the Transform type requires all generator types with a non-erased + // Transform to access the private constructor of the erased type, but sadly + // we cannot resonably restrict this to "T, non-void" without much more code + // or compiler warnings on some versions of clang, e.g. the one darwin uses. + template<typename, typename> + friend struct Generator; + + using promise_type = _generator::promise<T, Transform>; + + Generator(const Generator &) = delete; + Generator & operator=(const Generator &) = delete; + Generator(Generator &&) = default; + Generator & operator=(Generator &&) = default; + + /// If the coroutine held by the Generator has not finished, runs it until it + /// yields a value, throws any exception, or returns. If the coroutine yields + /// a value this value is passed to a persistent instance of `Transform` that + /// is held by the Generator, and the result of this call is returned. If the + /// coroutine throws an exception, or the Transform throws an exception while + /// processing an item, that exception is rethrown and the Generator will not + /// return any more non-`std::nullopt` values from next(). Once the contained + /// coroutine has completed or an exception has been thrown the Generator can + /// no longer return any valid values, only `std::nullopt`. Exceptions thrown + /// are thrown only once, further invocations of next() return `std::nullopt`. + /// + /// \returns `std::nullopt` if the coroutine has completed, or a value + std::optional<T> next() + { + return impl.next(); + } + + /// Type-erases the `Transform`. + /// + /// \return a new Generator with the `Transform` type-erased + Generator<T, void> decay() && + { + return Generator<T, void>(std::move(impl)); + } + + /// \copydoc decay() + operator Generator<T, void>() && + { + return std::move(*this).decay(); + } + +private: + _generator::GeneratorBase<T> impl; + + explicit Generator(_generator::GeneratorBase<T> b) : impl(std::move(b)) {} +}; + +} |