libutil: make generators iterable

range-for on generators sounds like a pretty good thing, right?

Change-Id: Ibed5e038c8dc50c918cf7c1f1aa70d822fb3efa2

lix/libutil/generator.hh

@@ -5,6 +5,7 @@
 
 #include <coroutine>
 #include <exception>
+#include <memory>
 #include <optional>
 #include <utility>
 #include <variant>
@@ -280,6 +281,71 @@ struct Generator
         return std::move(*this).decay();
     }
 
+    class iterator
+    {
+        // operator== must be const, but we need to call parent->next() to
+        // be able to check whether the sequence has ended. boldface sigh.
+        mutable Generator * parent = nullptr;
+        mutable std::shared_ptr<T> item = nullptr;
+
+        void step() const
+        {
+            auto next = parent->next();
+            if (!next) {
+                parent = nullptr;
+                item = nullptr;
+            } else if (!item) {
+                item = std::make_shared<T>(std::move(*next));
+            } else {
+                *item = std::move(*next);
+            }
+        }
+
+        void initializeIfNecessary() const
+        {
+            if (parent && !item) {
+                step();
+            }
+        }
+
+    public:
+        using iterator_category = std::input_iterator_tag;
+        using difference_type = void;
+        using value_type = T;
+        using reference = T &;
+        using pointer = T *;
+
+        iterator() = default;
+        explicit iterator(Generator & parent) : parent(&parent) {}
+
+        T * operator->() { return &**this; }
+        T & operator*()
+        {
+            initializeIfNecessary();
+            return *item;
+        }
+
+        iterator & operator++()
+        {
+            initializeIfNecessary();
+            if (parent) {
+                step();
+            }
+            return *this;
+        }
+
+        void operator++(int) { ++*this; }
+
+        bool operator==(const iterator & b) const
+        {
+            initializeIfNecessary();
+            return parent == nullptr && b.parent == nullptr;
+        }
+    };
+
+    iterator begin() { return iterator{*this}; }
+    iterator end() { return iterator{}; }
+
 private:
     _generator::GeneratorBase<T> impl;
 

tests/unit/libutil/generator.cc

@@ -216,4 +216,48 @@ TEST(Generator, transformThrows)
     }
 }
 
+TEST(Generator, iterators)
+{
+    auto g = []() -> Generator<int> {
+        for (auto i : {1, 2, 3, 4, 5, 6, 7, 8}) {
+            co_yield i;
+        }
+    }();
+
+    // begin() does not consume an item
+    {
+        auto it = g.begin();
+        ASSERT_EQ(g.next(), 1);
+    }
+
+    // operator* consumes only one item per advancement
+    {
+        auto it = g.begin();
+        ASSERT_EQ(*it, 2);
+        ASSERT_EQ(*it, 2);
+        ++it;
+        ASSERT_EQ(*it, 3);
+        ASSERT_EQ(*it, 3);
+    }
+
+    // not advancing an iterator consumes no items
+    ASSERT_EQ(g.next(), 4);
+
+    // operator++ on a fresh iterator consumes *two* items
+    {
+        auto it = g.begin();
+        ++it;
+        ASSERT_EQ(g.next(), 7);
+    }
+
+    // operator++ on last item reverts to end()
+    {
+        auto it = g.begin();
+        ASSERT_EQ(*it, 8);
+        ASSERT_NE(it, g.end());
+        ++it;
+        ASSERT_EQ(it, g.end());
+    }
+}
+
 }