feat: add flatten to iterator interface (#112)

Author: uncommon-nickname

include/rusty_iterators/flatten.hpp

@@ -0,0 +1,67 @@
+#pragma once
+
+#include "concepts.hpp"
+#include "interface.fwd.hpp"
+#include "peekable.hpp"
+
+#include <optional>
+
+namespace rusty_iterators::iterator
+{
+using concepts::Indexable;
+using interface::IterInterface;
+
+template <class T, class Other>
+    requires Indexable<T>
+class Flatten : public IterInterface<typename T::value_type, Flatten<T, Other>>
+{
+  public:
+    explicit Flatten(Other&& it) : it(std::forward<Peekable<T, Other>>(it.peekable())) {}
+
+    auto next() -> std::optional<typename T::value_type>;
+    [[nodiscard]] auto sizeHint() -> std::optional<size_t>;
+
+  private:
+    Peekable<T, Other> it;
+    size_t ptr             = 0;
+    std::optional<T> cache = std::nullopt;
+};
+} // namespace rusty_iterators::iterator
+
+template <class T, class Other>
+    requires rusty_iterators::concepts::Indexable<T>
+auto rusty_iterators::iterator::Flatten<T, Other>::next() -> std::optional<typename T::value_type>
+{
+    if (cache.has_value() && ptr < cache.value().size())
+    {
+        auto item = cache.value().at(ptr);
+        ptr += 1;
+        return std::move(item);
+    }
+    auto indexableItem = it.next();
+
+    if (!indexableItem.has_value())
+        return std::nullopt;
+
+    ptr   = 0;
+    cache = std::move(indexableItem);
+
+    return next();
+}
+
+template <class T, class Other>
+    requires rusty_iterators::concepts::Indexable<T>
+auto rusty_iterators::iterator::Flatten<T, Other>::sizeHint() -> std::optional<size_t>
+{
+    auto unfoldedSize = it.sizeHint();
+
+    if (!unfoldedSize.has_value())
+        return std::move(unfoldedSize);
+
+    auto peekedItem = it.peek();
+
+    if (!peekedItem.has_value())
+        return 0;
+
+    return unfoldedSize.value() * peekedItem.value().size();
+}

include/rusty_iterators/interface.hpp

@@ -6,6 +6,7 @@
 #include "enumerate.hpp"
 #include "filter.hpp"
 #include "filter_map.hpp"
+#include "flatten.hpp"
 #include "inspect.hpp"
 #include "interperse.hpp"
 #include "map.hpp"
@@ -47,6 +48,7 @@ using iterator::CycleType;
 using iterator::Enumerate;
 using iterator::Filter;
 using iterator::FilterMap;
+using iterator::Flatten;
 using iterator::Inspect;
 using iterator::Interperse;
 using iterator::Map;
@@ -113,6 +115,10 @@ class IterInterface
     [[nodiscard]] auto filterMap(Functor&& f)
         -> FilterMap<T, typename std::invoke_result_t<Functor, T>::value_type, Functor, Derived>;
 
+    template <class R = T>
+        requires Indexable<R>
+    [[nodiscard]] auto flatten() -> Flatten<R, Derived>;
+
     template <class B, class Functor>
         requires FoldFunctor<B, T, Functor>
     [[nodiscard]] auto fold(B&& init, Functor&& f) -> B;
@@ -241,7 +247,7 @@ auto rusty_iterators::interface::IterInterface<T, Derived>::collect() -> std::ve
     auto size       = sizeHintChecked();
 
     collection.reserve(size);
-    self().forEach([&collection](auto&& x) { collection.push_back(std::move(x)); });
+    self().forEach([&collection](auto x) { collection.push_back(x); });
 
     return std::move(collection);
 }
@@ -322,6 +328,14 @@ auto rusty_iterators::interface::IterInterface<T, Derived>::filterMap(Functor&&
         std::forward<Derived>(self()), std::forward<Functor>(f)};
 }
 
+template <class T, class Derived>
+template <class R>
+    requires rusty_iterators::concepts::Indexable<R>
+auto rusty_iterators::interface::IterInterface<T, Derived>::flatten() -> Flatten<R, Derived>
+{
+    return Flatten<R, Derived>{std::forward<Derived>(self())};
+}
+
 template <class T, class Derived>
 template <class B, class Functor>
     requires rusty_iterators::concepts::FoldFunctor<B, T, Functor>
@@ -353,7 +367,7 @@ auto rusty_iterators::interface::IterInterface<T, Derived>::forEach(Functor&& f)
 
     [[likely]] while (nextItem.has_value())
     {
-        func(std::move(nextItem.value()));
+        func(nextItem.value());
         nextItem = self().next();
     }
 }

include/rusty_iterators/iterator.hpp

@@ -16,6 +16,7 @@ using Item = std::reference_wrapper<const typename Container::value_type>;
 namespace rusty_iterators::iterator
 {
 using concepts::FoldFunctor;
+using concepts::Indexable;
 using concepts::Multiplyable;
 using concepts::Summable;
 
@@ -30,6 +31,10 @@ class LazyIterator : public interface::IterInterface<Item<Container>, LazyIterat
   public:
     explicit LazyIterator(Container& it) : ptr(it.begin()), end(it.end()) {}
 
+    template <class R = RawT>
+        requires Indexable<R>
+    [[nodiscard]] auto flatten() -> decltype(auto);
+
     auto next() -> std::optional<T>;
     [[nodiscard]] auto sizeHint() const -> std::optional<size_t>;
 
@@ -47,6 +52,15 @@ class LazyIterator : public interface::IterInterface<Item<Container>, LazyIterat
 };
 } // namespace rusty_iterators::iterator
 
+template <class Container>
+    requires std::ranges::range<Container>
+template <class R>
+    requires rusty_iterators::concepts::Indexable<R>
+auto rusty_iterators::iterator::LazyIterator<Container>::flatten() -> decltype(auto)
+{
+    return this->map([](auto x) { return x.get(); }).flatten();
+}
+
 template <class Container>
     requires std::ranges::range<Container>
 auto rusty_iterators::iterator::LazyIterator<Container>::next() -> std::optional<T>

tests/flatten.test.cpp

@@ -0,0 +1,66 @@
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include <rusty_iterators/iterator.hpp>
+
+using ::rusty_iterators::iterator::LazyIterator;
+using ::testing::ElementsAreArray;
+
+TEST(TestFlattenIterator, TestNextReturnsUnfolded)
+{
+    auto vec = std::vector<std::vector<int>>{
+        {1, 2},
+        {3, 4}
+    };
+    auto it = LazyIterator{vec}.flatten();
+
+    ASSERT_EQ(it.next(), 1);
+    ASSERT_EQ(it.next(), 2);
+    ASSERT_EQ(it.next(), 3);
+    ASSERT_EQ(it.next(), 4);
+    ASSERT_EQ(it.next(), std::nullopt);
+}
+
+TEST(TestFlattenIterator, TestCollectFlattened)
+{
+    auto vec = std::vector{1, 2, 3};
+    auto it  = LazyIterator{vec}.movingWindow(2).flatten();
+
+    EXPECT_THAT(it.collect(), ElementsAreArray({1, 2, 2, 3}));
+}
+
+TEST(TestFlattenIterator, TestDealingWithReferenceWrapper)
+{
+    auto vec = std::vector<std::vector<int>>{
+        {1, 2},
+        {3, 4}
+    };
+    auto it = LazyIterator{vec}.flatten();
+
+    EXPECT_THAT(it.collect(), ElementsAreArray({1, 2, 3, 4}));
+}
+
+TEST(TestFlattenIterator, TestSizeHintOnInfiniteIterator)
+{
+    auto vec = std::vector{1, 2, 3};
+    auto it  = LazyIterator{vec}.movingWindow(2).cycle().flatten();
+
+    ASSERT_EQ(it.sizeHint(), std::nullopt);
+}
+
+TEST(TestFlattenIterator, TestSizeHintOnEmptyIterator)
+{
+    auto vec = std::vector<int>{};
+    auto it  = LazyIterator{vec}.movingWindow(2).flatten();
+
+    ASSERT_EQ(it.sizeHint(), 0);
+}
+
+TEST(TestFlattenIterator, TestSizeHintOnNonEmptyIterator)
+{
+    auto vec = std::vector{1, 2, 3, 4};
+    auto it  = LazyIterator{vec}.movingWindow(2).flatten();
+
+    ASSERT_EQ(it.sizeHint(), 6);
+    ASSERT_EQ(it.count(), 6);
+}