libutil: generator type with on-yield value mapping

this will be the basis of non-boost coroutines in lix. anything that is
a boost coroutine *should* be representable with a Generator coroutine,
and many things that are not currently boost coroutines but behave much
like one (such as, notably, serializers) should be as well. this allows
us to greatly simplify many things that look like iteration but aren't.

Change-Id: I2cebcefa0148b631fb30df4c8cfa92167a407e34

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)) {}
+};
+
+}