[ADT] Fix zip iterator interface.

This commit provides `zip_{first,shortest}` with the standard member types and
methods expected of iterators (e.g., `difference_type`), in order for zip to be
used with other adaptors, such as `make_filter_range`.

Support for reverse iteration has also been added.

Differential Revision: https://reviews.llvm.org/D30246

llvm-svn: 296036

llvm/include/llvm/ADT/STLExtras.h

@@ -356,65 +356,120 @@ template <size_t... I> struct index_sequence;
 template <class... Ts> struct index_sequence_for;
 
 namespace detail {
-template <typename... Iters> class zip_first {
-public:
-  typedef std::input_iterator_tag iterator_category;
-  typedef std::tuple<decltype(*std::declval<Iters>())...> value_type;
+using std::declval;
+
+template <typename ZipType, typename... Iters>
+using zip_traits = iterator_facade_base<
+    ZipType, typename std::common_type<std::bidirectional_iterator_tag,
+                                       typename std::iterator_traits<
+                                           Iters>::iterator_category...>::type,
+    // ^ TODO: Implement random access methods.
+    std::tuple<decltype(*declval<Iters>())...>,
+    typename std::iterator_traits<typename std::tuple_element<
+        0, std::tuple<Iters...>>::type>::difference_type,
+    // ^ FIXME: This follows boost::make_zip_iterator's assumption that all
+    // inner iterators have the same difference_type. It would fail if, for
+    // instance, the second field's difference_type were non-numeric while the
+    // first is.
+    std::tuple<decltype(*declval<Iters>())...> *,
+    std::tuple<decltype(*declval<Iters>())...>>;
+
+template <typename ZipType, typename... Iters>
+struct zip_common : public zip_traits<ZipType, Iters...> {
+  using Base = zip_traits<ZipType, Iters...>;
+  using value_type = typename Base::value_type;
+
   std::tuple<Iters...> iterators;
 
-private:
-  template <size_t... Ns> value_type deres(index_sequence<Ns...>) {
+protected:
+  template <size_t... Ns> value_type deref(index_sequence<Ns...>) const {
     return value_type(*std::get<Ns>(iterators)...);
   }
 
-  template <size_t... Ns> decltype(iterators) tup_inc(index_sequence<Ns...>) {
+  template <size_t... Ns>
+  decltype(iterators) tup_inc(index_sequence<Ns...>) const {
     return std::tuple<Iters...>(std::next(std::get<Ns>(iterators))...);
   }
 
+  template <size_t... Ns>
+  decltype(iterators) tup_dec(index_sequence<Ns...>) const {
+    return std::tuple<Iters...>(std::prev(std::get<Ns>(iterators))...);
+  }
+
 public:
-  value_type operator*() { return deres(index_sequence_for<Iters...>{}); }
+  zip_common(Iters &&... ts) : iterators(std::forward<Iters>(ts)...) {}
+
+  value_type operator*() { return deref(index_sequence_for<Iters...>{}); }
+
+  const value_type operator*() const {
+    return deref(index_sequence_for<Iters...>{});
+  }
 
-  void operator++() { iterators = tup_inc(index_sequence_for<Iters...>{}); }
+  ZipType &operator++() {
+    iterators = tup_inc(index_sequence_for<Iters...>{});
+    return *reinterpret_cast<ZipType *>(this);
+  }
+
+  ZipType &operator--() {
+    static_assert(Base::IsBidirectional,
+                  "All inner iterators must be at least bidirectional.");
+    iterators = tup_dec(index_sequence_for<Iters...>{});
+    return *reinterpret_cast<ZipType *>(this);
+  }
+};
+
+template <typename... Iters>
+struct zip_first : public zip_common<zip_first<Iters...>, Iters...> {
+  using Base = zip_common<zip_first<Iters...>, Iters...>;
 
-  bool operator!=(const zip_first<Iters...> &other) const {
-    return std::get<0>(iterators) != std::get<0>(other.iterators);
+  bool operator==(const zip_first<Iters...> &other) const {
+    return std::get<0>(this->iterators) == std::get<0>(other.iterators);
   }
-  zip_first(Iters &&... ts) : iterators(std::forward<Iters>(ts)...) {}
+
+  zip_first(Iters &&... ts) : Base(std::forward<Iters>(ts)...) {}
 };
 
-template <typename... Iters> class zip_shortest : public zip_first<Iters...> {
+template <typename... Iters>
+class zip_shortest : public zip_common<zip_shortest<Iters...>, Iters...> {
   template <size_t... Ns>
-  bool test(const zip_first<Iters...> &other, index_sequence<Ns...>) const {
+  bool test(const zip_shortest<Iters...> &other, index_sequence<Ns...>) const {
     return all_of(std::initializer_list<bool>{std::get<Ns>(this->iterators) !=
                                               std::get<Ns>(other.iterators)...},
                   identity<bool>{});
   }
 
 public:
-  bool operator!=(const zip_first<Iters...> &other) const {
-    return test(other, index_sequence_for<Iters...>{});
+  using Base = zip_common<zip_shortest<Iters...>, Iters...>;
+
+  bool operator==(const zip_shortest<Iters...> &other) const {
+    return !test(other, index_sequence_for<Iters...>{});
   }
-  zip_shortest(Iters &&... ts)
-      : zip_first<Iters...>(std::forward<Iters>(ts)...) {}
+
+  zip_shortest(Iters &&... ts) : Base(std::forward<Iters>(ts)...) {}
 };
 
 template <template <typename...> class ItType, typename... Args> class zippy {
 public:
-  typedef ItType<decltype(std::begin(std::declval<Args>()))...> iterator;
+  using iterator = ItType<decltype(std::begin(std::declval<Args>()))...>;
+  using iterator_category = typename iterator::iterator_category;
+  using value_type = typename iterator::value_type;
+  using difference_type = typename iterator::difference_type;
+  using pointer = typename iterator::pointer;
+  using reference = typename iterator::reference;
 
 private:
   std::tuple<Args...> ts;
 
-  template <size_t... Ns> iterator begin_impl(index_sequence<Ns...>) {
+  template <size_t... Ns> iterator begin_impl(index_sequence<Ns...>) const {
     return iterator(std::begin(std::get<Ns>(ts))...);
   }
-  template <size_t... Ns> iterator end_impl(index_sequence<Ns...>) {
+  template <size_t... Ns> iterator end_impl(index_sequence<Ns...>) const {
     return iterator(std::end(std::get<Ns>(ts))...);
   }
 
 public:
-  iterator begin() { return begin_impl(index_sequence_for<Args...>{}); }
-  iterator end() { return end_impl(index_sequence_for<Args...>{}); }
+  iterator begin() const { return begin_impl(index_sequence_for<Args...>{}); }
+  iterator end() const { return end_impl(index_sequence_for<Args...>{}); }
   zippy(Args &&... ts_) : ts(std::forward<Args>(ts_)...) {}
 };
 } // End detail namespace

llvm/unittests/ADT/IteratorTest.cpp

@@ -7,9 +7,9 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ADT/iterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/iterator.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -35,14 +35,15 @@ static_assert(std::is_same<typename AdaptedIter::reference, Shadow<3>>::value,
               "");
 
 TEST(PointeeIteratorTest, Basic) {
-  int arr[4] = { 1, 2, 3, 4 };
+  int arr[4] = {1, 2, 3, 4};
   SmallVector<int *, 4> V;
   V.push_back(&arr[0]);
   V.push_back(&arr[1]);
   V.push_back(&arr[2]);
   V.push_back(&arr[3]);
 
-  typedef pointee_iterator<SmallVectorImpl<int *>::const_iterator> test_iterator;
+  typedef pointee_iterator<SmallVectorImpl<int *>::const_iterator>
+      test_iterator;
 
   test_iterator Begin, End;
   Begin = V.begin();
@@ -83,7 +84,8 @@ TEST(PointeeIteratorTest, SmartPointer) {
   V.push_back(make_unique<int>(4));
 
   typedef pointee_iterator<
-      SmallVectorImpl<std::unique_ptr<int>>::const_iterator> test_iterator;
+      SmallVectorImpl<std::unique_ptr<int>>::const_iterator>
+      test_iterator;
 
   test_iterator Begin, End;
   Begin = V.begin();
@@ -272,4 +274,51 @@ TEST(ZipIteratorTest, ZipFirstMutability) {
   }
 }
 
+TEST(ZipIteratorTest, Filter) {
+  using namespace std;
+  vector<unsigned> pi{3, 1, 4, 1, 5, 9};
+
+  unsigned iters = 0;
+  // pi is length 6, but the zip RHS is length 7.
+  auto zipped = zip_first(pi, vector<bool>{1, 1, 0, 1, 1, 1, 0});
+  for (auto tup : make_filter_range(
+           zipped, [](decltype(zipped)::value_type t) { return get<1>(t); })) {
+    EXPECT_EQ(get<0>(tup) & 0x01, get<1>(tup));
+    get<0>(tup) += 1;
+    iters += 1;
+  }
+
+  // Should have skipped pi[2].
+  EXPECT_EQ(iters, 5u);
+
+  // Ensure that in-place mutation works.
+  EXPECT_TRUE(all_of(pi, [](unsigned n) { return (n & 0x01) == 0; }));
+}
+
+TEST(ZipIteratorTest, Reverse) {
+  using namespace std;
+  vector<unsigned> ascending{0, 1, 2, 3, 4, 5};
+
+  auto zipped = zip_first(ascending, vector<bool>{0, 1, 0, 1, 0, 1});
+  unsigned last = 6;
+  for (auto tup : reverse(zipped)) {
+    // Check that this is in reverse.
+    EXPECT_LT(get<0>(tup), last);
+    last = get<0>(tup);
+    EXPECT_EQ(get<0>(tup) & 0x01, get<1>(tup));
+  }
+
+  auto odds = [](decltype(zipped)::value_type tup) { return get<1>(tup); };
+  last = 6;
+  for (auto tup : make_filter_range(reverse(zipped), odds)) {
+    EXPECT_LT(get<0>(tup), last);
+    last = get<0>(tup);
+    EXPECT_TRUE(get<0>(tup) & 0x01);
+    get<0>(tup) += 1;
+  }
+
+  // Ensure that in-place mutation works.
+  EXPECT_TRUE(all_of(ascending, [](unsigned n) { return (n & 0x01) == 0; }));
+}
+
 } // anonymous namespace