[PATCH] D138865: [ADT] Add `zip_equal` for iteratees of equal lengths

[Jakub Kuderski via Phabricator via llvm-commits]

kuhar created this revision.
kuhar added reviewers: mehdi_amini, Meinersbur, dblaikie, okkwon, mravishankar, benvanik.
Herald added a subscriber: hanchung.
Herald added a project: All.
kuhar requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D138865

Files:
  llvm/include/llvm/ADT/STLExtras.h
  llvm/unittests/ADT/IteratorTest.cpp


Index: llvm/unittests/ADT/IteratorTest.cpp
===================================================================
--- llvm/unittests/ADT/IteratorTest.cpp
+++ llvm/unittests/ADT/IteratorTest.cpp
@@ -416,6 +416,33 @@
   }
 }
 
+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};
Index: llvm/include/llvm/ADT/STLExtras.h
===================================================================
--- llvm/include/llvm/ADT/STLExtras.h
+++ llvm/include/llvm/ADT/STLExtras.h
@@ -720,12 +720,15 @@
                     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 @@
       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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D138865.478366.patch
Type: text/x-patch
Size: 3064 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20221128/0ddefe89/attachment.bin>