Add zip range adapter and other utility functions (#5425)

Last set of cherry picks for reverse iteration, this time enabling
utilities that the ReverseIterator test uses in upstream

---------

Co-authored-by: Tim Shen <[email protected]>
Co-authored-by: Duncan P. N. Exon Smith <[email protected]>
Co-authored-by: Justin Lebar <[email protected]>
Co-authored-by: Mehdi Amini <[email protected]>

include/llvm/ADT/STLExtras.h

@@ -343,6 +343,95 @@ make_filter_range(RangeT &&Range, PredicateT Pred) {
                     FilterIteratorT(std::end(std::forward<RangeT>(Range))));
 }
 
+// forward declarations required by zip_shortest/zip_first
+template <typename R, class UnaryPredicate>
+bool all_of(R &&range, UnaryPredicate &&P);
+
+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;
+  std::tuple<Iters...> iterators;
+
+private:
+  template <size_t... Ns> value_type deres(index_sequence<Ns...>) {
+    return value_type(*std::get<Ns>(iterators)...);
+  }
+
+  template <size_t... Ns> decltype(iterators) tup_inc(index_sequence<Ns...>) {
+    return std::tuple<Iters...>(std::next(std::get<Ns>(iterators))...);
+  }
+
+public:
+  value_type operator*() { return deres(index_sequence_for<Iters...>{}); }
+
+  void operator++() { iterators = tup_inc(index_sequence_for<Iters...>{}); }
+
+  bool operator!=(const zip_first<Iters...> &other) const {
+    return std::get<0>(iterators) != std::get<0>(other.iterators);
+  }
+  zip_first(Iters &&... ts) : iterators(std::forward<Iters>(ts)...) {}
+};
+
+template <typename... Iters> class zip_shortest : public zip_first<Iters...> {
+  template <size_t... Ns>
+  bool test(const zip_first<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...>{});
+  }
+  zip_shortest(Iters &&... ts)
+      : zip_first<Iters...>(std::forward<Iters>(ts)...) {}
+};
+
+template <template <typename...> class ItType, typename... Args> class zippy {
+public:
+  typedef ItType<decltype(std::begin(std::declval<Args>()))...> iterator;
+
+private:
+  std::tuple<Args...> ts;
+
+  template <size_t... Ns> iterator begin_impl(index_sequence<Ns...>) {
+    return iterator(std::begin(std::get<Ns>(ts))...);
+  }
+  template <size_t... Ns> iterator end_impl(index_sequence<Ns...>) {
+    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...>{}); }
+  zippy(Args &&... ts_) : ts(std::forward<Args>(ts_)...) {}
+};
+} // End detail namespace
+
+/// zip iterator for two or more iteratable types.
+template <typename T, typename U, typename... Args>
+detail::zippy<detail::zip_shortest, T, U, Args...> zip(T &&t, U &&u,
+                                                       Args &&... args) {
+  return detail::zippy<detail::zip_shortest, T, U, Args...>(
+      std::forward<T>(t), std::forward<U>(u), std::forward<Args>(args)...);
+}
+
+/// zip iterator that, for the sake of efficiency, assumes the first iteratee to
+/// be the shortest.
+template <typename T, typename U, typename... Args>
+detail::zippy<detail::zip_first, T, U, Args...> zip_first(T &&t, U &&u,
+                                                          Args &&... args) {
+  return detail::zippy<detail::zip_first, T, U, Args...>(
+      std::forward<T>(t), std::forward<U>(u), std::forward<Args>(args)...);
+}
+
 //===----------------------------------------------------------------------===//
 //     Extra additions to <utility>
 //===----------------------------------------------------------------------===//

include/llvm/ADT/iterator.h

@@ -233,6 +233,23 @@ struct pointee_iterator
   T &operator*() const { return **this->I; }
 };
 
+template <typename WrappedIteratorT,
+          typename T = decltype(&*std::declval<WrappedIteratorT>())>
+class pointer_iterator
+    : public iterator_adaptor_base<pointer_iterator<WrappedIteratorT>,
+                                   WrappedIteratorT, T> {
+  mutable T Ptr;
+
+public:
+  pointer_iterator() {}
+
+  explicit pointer_iterator(WrappedIteratorT u)
+      : pointer_iterator::iterator_adaptor_base(std::move(u)) {}
+
+  T &operator*() { return Ptr = &*this->I; }
+  const T &operator*() const { return Ptr = &*this->I; }
+};
+
 }
 
 #endif

unittests/ADT/IteratorTest.cpp

@@ -167,4 +167,91 @@ TEST(FilterIteratorTest, InputIterator) {
   EXPECT_EQ((SmallVector<int, 3>{1, 3, 5}), Actual);
 }
 
+TEST(PointerIterator, Basic) {
+  int A[] = {1, 2, 3, 4};
+  pointer_iterator<int *> Begin(std::begin(A)), End(std::end(A));
+  EXPECT_EQ(A, *Begin);
+  ++Begin;
+  EXPECT_EQ(A + 1, *Begin);
+  ++Begin;
+  EXPECT_EQ(A + 2, *Begin);
+  ++Begin;
+  EXPECT_EQ(A + 3, *Begin);
+  ++Begin;
+  EXPECT_EQ(Begin, End);
+}
+
+TEST(PointerIterator, Const) {
+  int A[] = {1, 2, 3, 4};
+  const pointer_iterator<int *> Begin(std::begin(A));
+  EXPECT_EQ(A, *Begin);
+  EXPECT_EQ(A + 1, std::next(*Begin, 1));
+  EXPECT_EQ(A + 2, std::next(*Begin, 2));
+  EXPECT_EQ(A + 3, std::next(*Begin, 3));
+  EXPECT_EQ(A + 4, std::next(*Begin, 4));
+}
+
+TEST(ZipIteratorTest, Basic) {
+  using namespace std;
+  const SmallVector<unsigned, 6> pi{3, 1, 4, 1, 5, 9};
+  SmallVector<bool, 6> odd{1, 1, 0, 1, 1, 1};
+  const char message[] = "yynyyy\0";
+
+  for (auto tup : zip(pi, odd, message)) {
+    EXPECT_EQ(get<0>(tup) & 0x01, get<1>(tup));
+    EXPECT_EQ(get<0>(tup) & 0x01 ? 'y' : 'n', get<2>(tup));
+  }
+
+  // note the rvalue
+  for (auto tup : zip(pi, SmallVector<bool, 0>{1, 1, 0, 1, 1})) {
+    EXPECT_EQ(get<0>(tup) & 0x01, get<1>(tup));
+  }
+}
+
+TEST(ZipIteratorTest, ZipFirstBasic) {
+  using namespace std;
+  const SmallVector<unsigned, 6> pi{3, 1, 4, 1, 5, 9};
+  unsigned iters = 0;
+
+  for (auto tup : zip_first(SmallVector<bool, 0>{1, 1, 0, 1}, pi)) {
+    EXPECT_EQ(get<0>(tup), get<1>(tup) & 0x01);
+    iters += 1;
+  }
+
+  EXPECT_EQ(iters, 4u);
+}
+
+TEST(ZipIteratorTest, Mutability) {
+  using namespace std;
+  const SmallVector<unsigned, 4> pi{3, 1, 4, 1, 5, 9};
+  char message[] = "hello zip\0";
+
+  for (auto tup : zip(pi, message, message)) {
+    EXPECT_EQ(get<1>(tup), get<2>(tup));
+    get<2>(tup) = get<0>(tup) & 0x01 ? 'y' : 'n';
+  }
+
+  // note the rvalue
+  for (auto tup : zip(message, "yynyyyzip\0")) {
+    EXPECT_EQ(get<0>(tup), get<1>(tup));
+  }
+}
+
+TEST(ZipIteratorTest, ZipFirstMutability) {
+  using namespace std;
+  vector<unsigned> pi{3, 1, 4, 1, 5, 9};
+  unsigned iters = 0;
+
+  for (auto tup : zip_first(SmallVector<bool, 0>{1, 1, 0, 1}, pi)) {
+    get<1>(tup) = get<0>(tup);
+    iters += 1;
+  }
+
+  EXPECT_EQ(iters, 4u);
+
+  for (auto tup : zip_first(SmallVector<bool, 0>{1, 1, 0, 1}, pi)) {
+    EXPECT_EQ(get<0>(tup), get<1>(tup));
+  }
+}
+
 } // anonymous namespace