From 21d434d997db58afcf03f72c5a34a56d2d46f6d1 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <kubak@google.com>
Date: Tue, 29 Nov 2022 19:56:23 -0500
Subject: [PATCH] [ADT] Add `zip_equal` for iteratees of equal lengths

Add a new version of `zip` that assumes that all iteratees have equal
lengths. The difference compared to `zip_first` is that `zip_equal`
checks this assumption in builds with assertions enabled.

This will allow us to clearly express the intent when working with
equally-sized ranges without having to write this assertion manually.

This is similar to Python's `zip(..., equal=True)` [1] or
`more_itertools.zip_equal` [2].

I saw this first suggested by @benvanik.

[1] https://peps.python.org/pep-0618/
[2] https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.zip_equal

Reviewed By: dblaikie

Differential Revision: https://reviews.llvm.org/D138865
---
 llvm/include/llvm/ADT/STLExtras.h   | 19 ++++++++++++++++++-
 llvm/unittests/ADT/IteratorTest.cpp | 27 +++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h
index 8e290cdb71053..11ef2f99736e7 100644
--- a/llvm/include/llvm/ADT/STLExtras.h
+++ b/llvm/include/llvm/ADT/STLExtras.h
@@ -720,12 +720,15 @@ make_early_inc_range(RangeT &&Range) {
                     EarlyIncIteratorT(std::end(std::forward<RangeT>(Range))));
 }
 
-// forward declarations required by zip_shortest/zip_first/zip_longest
+// Forward declarations required by zip_shortest/zip_equal/zip_first/zip_longest
 template <typename R, typename UnaryPredicate>
 bool all_of(R &&range, UnaryPredicate P);
+
 template <typename R, typename UnaryPredicate>
 bool any_of(R &&range, UnaryPredicate P);
 
+template <typename T> bool all_equal(std::initializer_list<T> Values);
+
 namespace detail {
 
 using std::declval;
@@ -879,6 +882,20 @@ detail::zippy<detail::zip_shortest, T, U, Args...> zip(T &&t, U &&u,
       std::forward<T>(t), std::forward<U>(u), std::forward<Args>(args)...);
 }
 
+/// zip iterator that assumes that all iteratees have the same length.
+/// In builds with assertions on, this assumption is checked before the
+/// iteration starts.
+template <typename T, typename U, typename... Args>
+detail::zippy<detail::zip_first, T, U, Args...> zip_equal(T &&t, U &&u,
+                                                          Args &&...args) {
+  assert(all_equal({std::distance(adl_begin(t), adl_end(t)),
+                    std::distance(adl_begin(u), adl_end(u)),
+                    std::distance(adl_begin(args), adl_end(args))...}) &&
+         "Iteratees do not have equal length");
+  return detail::zippy<detail::zip_first, 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. Iteration continues until the end of the first iteratee is
 /// reached. In builds with assertions on, we check that the assumption about
diff --git a/llvm/unittests/ADT/IteratorTest.cpp b/llvm/unittests/ADT/IteratorTest.cpp
index 4d68d0faff53c..14fbcd5cf84c0 100644
--- a/llvm/unittests/ADT/IteratorTest.cpp
+++ b/llvm/unittests/ADT/IteratorTest.cpp
@@ -416,6 +416,33 @@ TEST(ZipIteratorTest, Basic) {
   }
 }
 
+TEST(ZipIteratorTest, ZipEqualBasic) {
+  const SmallVector<unsigned, 6> pi = {3, 1, 4, 1, 5, 8};
+  const SmallVector<bool, 6> vals = {1, 1, 0, 1, 1, 0};
+  unsigned iters = 0;
+
+  for (auto [lhs, rhs] : zip_equal(vals, pi)) {
+    EXPECT_EQ(lhs, rhs & 0x01);
+    ++iters;
+  }
+
+  EXPECT_EQ(iters, 6u);
+}
+
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+// Check that an assertion is triggered when ranges passed to `zip_equal` differ
+// in length.
+TEST(ZipIteratorTest, ZipEqualNotEqual) {
+  const SmallVector<unsigned, 6> pi = {3, 1, 4, 1, 5, 8};
+  const SmallVector<bool, 2> vals = {1, 1};
+
+  EXPECT_DEATH(zip_equal(pi, vals), "Iteratees do not have equal length");
+  EXPECT_DEATH(zip_equal(vals, pi), "Iteratees do not have equal length");
+  EXPECT_DEATH(zip_equal(pi, pi, vals), "Iteratees do not have equal length");
+  EXPECT_DEATH(zip_equal(vals, vals, pi), "Iteratees do not have equal length");
+}
+#endif
+
 TEST(ZipIteratorTest, ZipFirstBasic) {
   using namespace std;
   const SmallVector<unsigned, 6> pi{3, 1, 4, 1, 5, 9};