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
2024-06-30 17:23:41 +02:00
parent 5dc85e8b72
commit 11307e04d1
5 changed files with 476 additions and 0 deletions
@@ -142,6 +142,16 @@ else
  cpp_pch = []
 endif

+# gcc 12 is known to miscompile some coroutine-based code quite horribly,
+# causing (among other things) copies of move-only objects and the double
+# frees one would expect when the objects are unique_ptrs. these problems
+# often show up as memory corruption when nesting generators (since we do
+# treat generators like owned memory) and will cause inexplicable crashs.
+assert(
+  cxx.get_id() != 'gcc' or cxx.version().version_compare('>=13'),
+  'GCC 12 and earlier are known to miscompile lix coroutines, use GCC 13 or clang.'
+)
+

 # Translate some historical and Mesony CPU names to Lixy CPU names.
 # FIXME(Qyriad): the 32-bit x86 code is not tested right now, because cross compilation for Lix
@@ -0,0 +1,250 @@
+#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
+{
+    std::variant<T, link<T>, failure, finished> value{};
+    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;
+
+    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()};
+    }
+
+    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 {};
+    }
+
+    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 {};
+    }
+
+    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()
+    {
+        while (active.handle) {
+            active.handle.resume();
+            auto & p = *active.state;
+            auto result = std::visit(
+                overloaded{
+                    [&](finished) -> std::optional<T> {
+                        active = p.parent;
+                        return {};
+                    },
+                    [&](link<T> & inner) -> std::optional<T> {
+                        auto base = inner.state;
+                        while (base->parent.handle) {
+                            base = base->parent.state;
+                        }
+                        base->parent = active;
+                        active = inner;
+                        return {};
+                    },
+                    [&](T & value) -> std::optional<T> { return std::move(value); },
+                    [&](failure & f) -> std::optional<T> {
+                        active = {};
+                        std::rethrow_exception(f.e);
+                    },
+                },
+                p.value
+            );
+            if (result) {
+                return std::move(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);
+    }
+};
+
+}
+
+/// 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;
+    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 and 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)) {}
+};
+
+}
@@ -72,6 +72,7 @@ libutil_headers = files(
  'file-system.hh',
  'finally.hh',
  'fmt.hh',
+  'generator.hh',
  'git.hh',
  'hash.hh',
  'hilite.hh',
@@ -0,0 +1,214 @@
+#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_EQ(g.next(), 1);
+    ASSERT_EQ(g.next(), 2);
+    ASSERT_FALSE(g.next().has_value());
+}
+
+TEST(Generator, returns)
+{
+    {
+        auto g = []() -> Generator<int> { co_return; }();
+
+        ASSERT_FALSE(g.next().has_value());
+    }
+    {
+        auto g = []() -> Generator<int> {
+            co_yield 1;
+            co_yield []() -> Generator<int> { co_return; }();
+            co_yield 2;
+            co_yield []() -> Generator<int> { co_yield 10; }();
+            co_yield 3;
+            (void) "dummy statement to force some more execution";
+        }();
+
+        ASSERT_EQ(g.next(), 1);
+        ASSERT_EQ(g.next(), 2);
+        ASSERT_EQ(g.next(), 10);
+        ASSERT_EQ(g.next(), 3);
+        ASSERT_FALSE(g.next().has_value());
+    }
+}
+
+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.next().value();
+        co_yield std::move(g2);
+        co_yield 2;
+    }();
+
+    ASSERT_EQ(g.next(), 1);
+    ASSERT_EQ(g.next(), 9);
+    ASSERT_EQ(g.next(), 99);
+    ASSERT_EQ(g.next(), 100);
+    ASSERT_EQ(g.next(), 2000);
+    ASSERT_EQ(g.next(), 2001);
+    ASSERT_EQ(g.next(), 1001);
+    ASSERT_EQ(g.next(), 2);
+    ASSERT_FALSE(g.next().has_value());
+}
+
+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_EQ(g.next(), 1);
+    ASSERT_EQ(g.next(), 9);
+    ASSERT_THROW(g.next(), int);
+}
+
+TEST(Generator, exception)
+{
+    {
+        auto g = []() -> Generator<int> {
+            co_yield 1;
+            throw 1;
+        }();
+
+        ASSERT_EQ(g.next(), 1);
+        ASSERT_THROW(g.next(), int);
+        ASSERT_FALSE(g.next().has_value());
+    }
+    {
+        auto g = []() -> Generator<int> {
+            throw 1;
+            co_return;
+        }();
+
+        ASSERT_THROW(g.next(), int);
+        ASSERT_FALSE(g.next().has_value());
+    }
+}
+
+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 (auto i = g.next()) {
+                co_yield *i;
+            }
+        }(std::move(inner));
+    }
+};
+}
+
+TEST(Generator, transform)
+{
+    auto g = []() -> Generator<std::pair<uint32_t, int>, Transform> {
+        co_yield int32_t(-1);
+        co_yield "";
+        co_yield []() -> Generator<int> { co_yield 7; }();
+        co_yield 20;
+    }();
+
+    ASSERT_EQ(g.next(), (std::pair<unsigned, int>{4294967295, 0}));
+    ASSERT_EQ(g.next(), (std::pair<unsigned, int>{9, 0}));
+    ASSERT_EQ(g.next(), (std::pair<unsigned, int>{19, 1}));
+    ASSERT_EQ(g.next(), (std::pair<unsigned, int>{7, 0}));
+    ASSERT_EQ(g.next(), (std::pair<unsigned, int>{20, 1}));
+    ASSERT_FALSE(g.next().has_value());
+}
+
+namespace {
+struct ThrowTransform
+{
+    int operator()(int x)
+    {
+        return x;
+    }
+
+    int operator()(bool)
+    {
+        throw 2;
+    }
+
+    Generator<int, void> operator()(Generator<int> && inner)
+    {
+        throw false;
+    }
+};
+}
+
+TEST(Generator, transformThrows)
+{
+    {
+        auto g = []() -> Generator<int, ThrowTransform> {
+            co_yield 1;
+            co_yield false;
+            co_yield 2;
+        }();
+
+        ASSERT_EQ(g.next(), 1);
+        ASSERT_THROW(g.next(), int);
+        ASSERT_FALSE(g.next().has_value());
+    }
+    {
+        auto g = []() -> Generator<int, ThrowTransform> {
+            co_yield 1;
+            co_yield []() -> Generator<int> {
+                co_yield 2;
+            }();
+            co_yield 3;
+        }();
+
+        ASSERT_EQ(g.next(), 1);
+        ASSERT_THROW(g.next(), bool);
+        ASSERT_FALSE(g.next().has_value());
+    }
+}
+
+}
@@ -40,6 +40,7 @@ libutil_tests_sources = files(
  'libutil/compression.cc',
  'libutil/config.cc',
  'libutil/escape-string.cc',
+  'libutil/generator.cc',
  'libutil/git.cc',
  'libutil/hash.cc',
  'libutil/hilite.cc',