diff --git a/lix/libutil/generator.hh b/lix/libutil/generator.hh index e6d674466..e81a7f7c8 100644 --- a/lix/libutil/generator.hh +++ b/lix/libutil/generator.hh @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -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 item = nullptr; + + void step() const + { + auto next = parent->next(); + if (!next) { + parent = nullptr; + item = nullptr; + } else if (!item) { + item = std::make_shared(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 impl; diff --git a/tests/unit/libutil/generator.cc b/tests/unit/libutil/generator.cc index 4b8676a63..b08cb29bb 100644 --- a/tests/unit/libutil/generator.cc +++ b/tests/unit/libutil/generator.cc @@ -216,4 +216,48 @@ TEST(Generator, transformThrows) } } +TEST(Generator, iterators) +{ + auto g = []() -> Generator { + 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()); + } +} + }