Move the FlattenIterator in SPARTA

Summary:
Mariana Trench implements a `FlattenIterator` which can be used to get an iterator on a container of container.
This can be useful to other tools, so let's move it in sparta. This will also be used in a follow-up diff.

Reviewed By: anwesht, arnaudvenet

Differential Revision: D49607366

fbshipit-source-id: d438cb6208f004cdb40ff0fd4c760d6658736275

include/sparta/FlattenIterator.h

@@ -0,0 +1,132 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <iterator>
+#include <optional>
+#include <type_traits>
+
+namespace sparta {
+
+namespace fi_impl {
+
+template <typename T>
+struct Range {
+  T begin;
+  T end;
+};
+
+template <typename OuterIterator>
+struct FlattenDereference {
+  using Reference = typename std::iterator_traits<OuterIterator>::reference;
+  using InnerIterator = decltype(std::declval<Reference>().begin());
+
+  static InnerIterator begin(Reference reference) { return reference.begin(); }
+
+  static InnerIterator end(Reference reference) { return reference.end(); }
+};
+
+template <typename OuterIterator>
+struct FlattenConstDereference {
+  using Reference = typename std::iterator_traits<OuterIterator>::reference;
+  using InnerIterator = decltype(std::declval<Reference>().cbegin());
+
+  static InnerIterator begin(Reference reference) { return reference.cbegin(); }
+
+  static InnerIterator end(Reference reference) { return reference.cend(); }
+};
+
+} // namespace fi_impl
+
+/**
+ * A flattening iterator that iterates on a container of containers.
+ *
+ * For instance, this can be used to treat a `std::vector<std::vector<T>>` as
+ * a single list of `T`.
+ */
+template <typename OuterIterator,
+          typename InnerIterator,
+          typename Dereference = fi_impl::FlattenDereference<OuterIterator>>
+class FlattenIterator {
+ public:
+  using OuterReference =
+      typename std::iterator_traits<OuterIterator>::reference;
+
+  static_assert(std::is_same_v<
+                decltype(Dereference::begin(std::declval<OuterReference>())),
+                InnerIterator>);
+  static_assert(
+      std::is_same_v<decltype(Dereference::end(std::declval<OuterReference>())),
+                     InnerIterator>);
+
+  // C++ iterator concept member types
+  using iterator_category = std::forward_iterator_tag;
+  using value_type = typename std::iterator_traits<InnerIterator>::value_type;
+  using difference_type =
+      typename std::iterator_traits<OuterIterator>::difference_type;
+  using pointer = typename std::iterator_traits<InnerIterator>::pointer;
+  using reference = typename std::iterator_traits<InnerIterator>::reference;
+
+  explicit FlattenIterator(OuterIterator begin, OuterIterator end)
+      : m_outer(
+            fi_impl::Range<OuterIterator>{std::move(begin), std::move(end)}),
+        m_inner(std::nullopt) {
+    if (m_outer.begin == m_outer.end) {
+      return;
+    }
+    m_inner = fi_impl::Range<InnerIterator>{Dereference::begin(*m_outer.begin),
+                                            Dereference::end(*m_outer.begin)};
+    advance_empty();
+  }
+
+  FlattenIterator& operator++() {
+    ++m_inner->begin;
+    advance_empty();
+    return *this;
+  }
+
+  FlattenIterator operator++(int) {
+    FlattenIterator result = *this;
+    ++(*this);
+    return result;
+  }
+
+  bool operator==(const FlattenIterator& other) const {
+    return m_outer.begin == other.m_outer.begin &&
+           ((!m_inner.has_value() && !other.m_inner.has_value()) ||
+            (m_inner.has_value() && other.m_inner.has_value() &&
+             m_inner->begin == other.m_inner->begin));
+  }
+
+  bool operator!=(const FlattenIterator& other) const {
+    return !(*this == other);
+  }
+
+  reference operator*() { return *m_inner->begin; }
+
+ private:
+  /* Advance the iterator until we find an element. */
+  void advance_empty() {
+    while (m_inner->begin == m_inner->end) {
+      ++m_outer.begin;
+      if (m_outer.begin == m_outer.end) {
+        m_inner = std::nullopt;
+        return;
+      } else {
+        m_inner =
+            fi_impl::Range<InnerIterator>{Dereference::begin(*m_outer.begin),
+                                          Dereference::end(*m_outer.begin)};
+      }
+    }
+  }
+
+  fi_impl::Range<OuterIterator> m_outer;
+  std::optional<fi_impl::Range<InnerIterator>> m_inner;
+};
+
+} // namespace sparta

test/FlattenIteratorTest.cpp

@@ -0,0 +1,200 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <list>
+#include <map>
+#include <vector>
+
+#include <gmock/gmock.h>
+
+#include <sparta/FlattenIterator.h>
+
+using namespace sparta;
+
+namespace {
+
+template <typename Iterator>
+std::vector<typename std::iterator_traits<Iterator>::value_type> collect(
+    Iterator begin, Iterator end) {
+  std::vector<typename std::iterator_traits<Iterator>::value_type> result;
+  for (; begin != end; ++begin) {
+    result.push_back(*begin);
+  }
+  return result;
+}
+
+} // namespace
+
+TEST(FlattenIteratorTest, VectorVectorInt) {
+  using Vector = std::vector<int>;
+  using VectorVector = std::vector<std::vector<int>>;
+  using Iterator = FlattenIterator<
+      /* OuterIterator */ std::vector<std::vector<int>>::iterator,
+      /* InnerIterator */ std::vector<int>::iterator>;
+
+  VectorVector container = {};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            Vector{});
+
+  container = {{1}, {2, 3}, {4, 5, 6}};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{}, {1}, {}, {2, 3}, {}, {4, 5, 6}, {}};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{1}, {}, {2, 3}, {}, {4, 5}, {}, {6}};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+}
+
+TEST(FlattenIteratorTest, ListVectorInt) {
+  using Vector = std::vector<int>;
+  using ListVector = std::list<std::vector<int>>;
+  using Iterator = FlattenIterator<
+      /* OuterIterator */ std::list<std::vector<int>>::iterator,
+      /* InnerIterator */ std::vector<int>::iterator>;
+
+  ListVector container = {};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            Vector{});
+
+  container = {{1}, {2, 3}, {4, 5, 6}};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{}, {1}, {}, {2, 3}, {}, {4, 5, 6}, {}};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{1}, {}, {2, 3}, {}, {4, 5}, {}, {6}};
+  EXPECT_EQ(collect(Iterator(container.begin(), container.end()),
+                    Iterator(container.end(), container.end())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+}
+
+TEST(FlattenIteratorTest, ConstVectorVectorInt) {
+  using Vector = std::vector<int>;
+  using VectorVector = std::vector<std::vector<int>>;
+  using Iterator = FlattenIterator<
+      /* OuterIterator */ std::vector<std::vector<int>>::const_iterator,
+      /* InnerIterator */ std::vector<int>::const_iterator>;
+
+  VectorVector container = {};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            Vector{});
+
+  container = {{1}, {2, 3}, {4, 5, 6}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{}, {1}, {}, {2, 3}, {}, {4, 5, 6}, {}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{1}, {}, {2, 3}, {}, {4, 5}, {}, {6}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+}
+
+TEST(FlattenIteratorTest, MapVectorInt) {
+  using Vector = std::vector<int>;
+  using MapVector = std::map<int, std::vector<int>>;
+
+  struct Dereference {
+    static std::vector<int>::const_iterator begin(
+        const std::pair<const int, std::vector<int>>& p) {
+      return p.second.cbegin();
+    }
+    static std::vector<int>::const_iterator end(
+        const std::pair<const int, std::vector<int>>& p) {
+      return p.second.cend();
+    }
+  };
+
+  using Iterator = FlattenIterator<
+      /* OuterIterator */ std::map<int, std::vector<int>>::const_iterator,
+      /* InnerIterator */ std::vector<int>::const_iterator,
+      Dereference>;
+
+  MapVector container = {};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            Vector{});
+
+  container = {{0, {1}}, {1, {2, 3}}, {3, {4, 5, 6}}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{0, {}}, {1, {1}},       {2, {}}, {3, {2, 3}},
+               {4, {}}, {5, {4, 5, 6}}, {6, {}}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{0, {1}},    {1, {}}, {2, {2, 3}}, {3, {}},
+               {4, {4, 5}}, {5, {}}, {6, {6}}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+}
+
+TEST(FlattenIteratorTest, MapListInt) {
+  using Vector = std::vector<int>;
+  using MapVector = std::map<int, std::list<int>>;
+
+  struct Dereference {
+    static std::list<int>::const_iterator begin(
+        const std::pair<const int, std::list<int>>& p) {
+      return p.second.cbegin();
+    }
+    static std::list<int>::const_iterator end(
+        const std::pair<const int, std::list<int>>& p) {
+      return p.second.cend();
+    }
+  };
+
+  using Iterator = FlattenIterator<
+      /* OuterIterator */ std::map<int, std::list<int>>::const_iterator,
+      /* InnerIterator */ std::list<int>::const_iterator,
+      Dereference>;
+
+  MapVector container = {};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            Vector{});
+
+  container = {{0, {1}}, {1, {2, 3}}, {3, {4, 5, 6}}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{0, {}}, {1, {1}},       {2, {}}, {3, {2, 3}},
+               {4, {}}, {5, {4, 5, 6}}, {6, {}}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+
+  container = {{0, {1}},    {1, {}}, {2, {2, 3}}, {3, {}},
+               {4, {4, 5}}, {5, {}}, {6, {6}}};
+  EXPECT_EQ(collect(Iterator(container.cbegin(), container.cend()),
+                    Iterator(container.cend(), container.cend())),
+            (Vector{1, 2, 3, 4, 5, 6}));
+}