[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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@296036 91177308-0d34-0410-b5e6-96231b3b80d8

1 files changed, 53 insertions, 4 deletions

diff --git a/unittests/ADT/IteratorTest.cpp b/unittests/ADT/IteratorTest.cpp
index a8d5b33a0b4..0edb98b9343 100644
--- a/unittests/ADT/IteratorTest.cpp
+++ b/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