libutil: basic generator type with mapping

Change-Id: I2cebcefa0148b631fb30df4c8cfa92167a407e34

src/libutil/generator.hh

@@ -0,0 +1,181 @@
+#pragma once
+
+#include "overloaded.hh"
+
+#include <coroutine>
+#include <optional>
+#include <utility>
+#include <variant>
+
+namespace nix {
+
+template<typename T, typename Transform = std::identity>
+struct Generator : private Generator<T, void>
+{
+    struct promise_type;
+    using handle_type = std::coroutine_handle<promise_type>;
+
+    explicit Generator(handle_type h) : Generator<T, void>{h, h.promise().state} {}
+
+    using Generator<T, void>::operator bool;
+    using Generator<T, void>::operator();
+
+    operator Generator<T, void> &() &
+    {
+        return *this;
+    }
+    operator Generator<T, void>() &&
+    {
+        return std::move(*this);
+    }
+};
+
+template<typename T>
+struct Generator<T, void>
+{
+    template<typename T2, typename Transform>
+    friend struct Generator<T2, Transform>::promise_type;
+
+    struct promise_state;
+
+    struct _link
+    {
+        std::coroutine_handle<> handle{};
+        promise_state * state{};
+    };
+
+    struct promise_state
+    {
+        std::variant<_link, T> value{};
+        std::exception_ptr exception{};
+        _link parent{};
+    };
+
+    // 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,
+    // which mostly means move-only.
+    Generator(Generator && other)
+    {
+        swap(other);
+    }
+
+    Generator & operator=(Generator && other)
+    {
+        Generator(std::move(other)).swap(*this);
+        return *this;
+    }
+
+    ~Generator()
+    {
+        if (h) {
+            h.destroy();
+        }
+    }
+
+    explicit operator bool()
+    {
+        return ensure();
+    }
+
+    T operator()()
+    {
+        ensure();
+        auto result = std::move(*current);
+        current = nullptr;
+        return result;
+    }
+
+protected:
+    std::coroutine_handle<> h{};
+    _link active{};
+    T * current{};
+
+    Generator(std::coroutine_handle<> h, promise_state & state) : h(h), active(h, &state) {}
+
+    void swap(Generator & other)
+    {
+        std::swap(h, other.h);
+        std::swap(active, other.active);
+        std::swap(current, other.current);
+    }
+
+    bool ensure()
+    {
+        while (!current && active.handle) {
+            active.handle.resume();
+            auto & p = *active.state;
+            if (p.exception) {
+                std::rethrow_exception(p.exception);
+            } else if (active.handle.done()) {
+                active = p.parent;
+            } else {
+                std::visit(
+                    overloaded{
+                        [&](_link & inner) {
+                            auto base = inner.state;
+                            while (base->parent.handle) {
+                                base = base->parent.state;
+                            }
+                            base->parent = active;
+                            active = inner;
+                        },
+                        [&](T & value) { current = &value; },
+                    },
+                    p.value
+                );
+            }
+        }
+        return current;
+    }
+};
+
+template<typename T, typename Transform>
+struct Generator<T, Transform>::promise_type
+{
+    Generator<T, void>::promise_state state;
+    Transform convert;
+    std::optional<Generator<T, void>> inner;
+
+    Generator get_return_object()
+    {
+        return Generator(handle_type::from_promise(*this));
+    }
+    std::suspend_always initial_suspend()
+    {
+        return {};
+    }
+    std::suspend_always final_suspend() noexcept
+    {
+        return {};
+    }
+    void unhandled_exception()
+    {
+        state.exception = std::current_exception();
+    }
+
+    template<typename From>
+        requires requires(Transform t, From && f) {
+            {
+                t(std::forward<From>(f))
+            } -> std::convertible_to<T>;
+        }
+    std::suspend_always yield_value(From && from)
+    {
+        state.value = convert(std::forward<From>(from));
+        return {};
+    }
+
+    template<typename From>
+        requires requires(Transform t, From f) { static_cast<Generator<T, void>>(t(std::move(f))); }
+    std::suspend_always yield_value(From from)
+    {
+        inner = static_cast<Generator<T, void>>(convert(std::move(from)));
+        state.value = inner->active;
+        return {};
+    }
+
+    void return_void() {}
+};
+
+}

tests/unit/libutil/generator.cc

@@ -0,0 +1,141 @@
+#include "generator.hh"
+
+#include <concepts>
+#include <cstdint>
+#include <gtest/gtest.h>
+
+namespace nix {
+
+TEST(Generator, yields)
+{
+    auto g = []() -> Generator<int> {
+        co_yield 1;
+        co_yield 2;
+    }();
+
+    ASSERT_TRUE(bool(g));
+    ASSERT_EQ(g(), 1);
+    ASSERT_EQ(g(), 2);
+    ASSERT_FALSE(bool(g));
+}
+
+TEST(Generator, nests)
+{
+    auto g = []() -> Generator<int> {
+        co_yield 1;
+        co_yield []() -> Generator<int> {
+            co_yield 9;
+            co_yield []() -> Generator<int> {
+                co_yield 99;
+                co_yield 100;
+            }();
+        }();
+
+        auto g2 = []() -> Generator<int> {
+            co_yield []() -> Generator<int> {
+                co_yield 2000;
+                co_yield 2001;
+            }();
+            co_yield 1001;
+        }();
+
+        co_yield g2();
+        co_yield std::move(g2);
+        co_yield 2;
+    }();
+
+    ASSERT_TRUE(bool(g));
+    ASSERT_EQ(g(), 1);
+    ASSERT_EQ(g(), 9);
+    ASSERT_EQ(g(), 99);
+    ASSERT_EQ(g(), 100);
+    ASSERT_EQ(g(), 2000);
+    ASSERT_EQ(g(), 2001);
+    ASSERT_EQ(g(), 1001);
+    ASSERT_EQ(g(), 2);
+    ASSERT_FALSE(bool(g));
+}
+
+TEST(Generator, nestsExceptions)
+{
+    auto g = []() -> Generator<int> {
+        co_yield 1;
+        co_yield []() -> Generator<int> {
+            co_yield 9;
+            throw 1;
+            co_yield 10;
+        }();
+        co_yield 2;
+    }();
+
+    ASSERT_TRUE(bool(g));
+    ASSERT_EQ(g(), 1);
+    ASSERT_EQ(g(), 9);
+    ASSERT_THROW(g(), int);
+}
+
+TEST(Generator, exception)
+{
+    {
+        auto g = []() -> Generator<int> {
+            throw 1;
+            co_return;
+        }();
+
+        ASSERT_THROW(void(bool(g)), int);
+    }
+    {
+        auto g = []() -> Generator<int> {
+            throw 1;
+            co_return;
+        }();
+
+        ASSERT_THROW(g(), int);
+    }
+}
+
+namespace {
+struct Transform
+{
+    int state = 0;
+
+    std::pair<uint32_t, int> operator()(std::integral auto x)
+    {
+        return {x, state++};
+    }
+
+    Generator<std::pair<uint32_t, int>, Transform> operator()(const char *)
+    {
+        co_yield 9;
+        co_yield 19;
+    }
+
+    Generator<std::pair<uint32_t, int>, Transform> operator()(Generator<int> && inner)
+    {
+        return [](auto g) mutable -> Generator<std::pair<uint32_t, int>, Transform> {
+            while (g) {
+                co_yield g();
+            }
+        }(std::move(inner));
+    }
+};
+}
+
+TEST(Generator, transform)
+{
+    auto g = []() -> Generator<std::pair<uint32_t, int>, Transform> {
+        co_yield int32_t(-1);
+        co_yield "";
+        std::cerr << "1\n";
+        co_yield []() -> Generator<int> { co_yield 7; }();
+        co_yield 20;
+    }();
+
+    ASSERT_EQ(g(), (std::pair<unsigned, int>{4294967295, 0}));
+    ASSERT_EQ(g(), (std::pair<unsigned, int>{9, 0}));
+    ASSERT_EQ(g(), (std::pair<unsigned, int>{19, 1}));
+    ASSERT_EQ(g(), (std::pair<unsigned, int>{7, 0}));
+    ASSERT_EQ(g(), (std::pair<unsigned, int>{20, 1}));
+}
+
+}