[ADT] Add `zip_equal` for iteratees of equal lengths

[kuhar]

Summary

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

[dblaikie]

Sounds good :)

llvm/include/llvm/ADT/STLExtras.h

 iterator_range<early_inc_iterator_impl<detail::IterOfRange<RangeT>>>
 make_early_inc_range(RangeT &&Range) {
   using EarlyIncIteratorT =
       early_inc_iterator_impl<detail::IterOfRange<RangeT>>;
   return make_range(EarlyIncIteratorT(std::begin(std::forward<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;
 
 // We have to alias this since inlining the actual type at the usage site
 // in the parameter list of iterator_facade_base<> below ICEs MSVC 2017.
 template<typename... Iters> struct ZipTupleType {
   using type = std::tuple<decltype(*declval<Iters>())...>;

 /// end of the *shortest* iteratee is reached.
 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 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.
 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)...);

llvm/unittests/ADT/IteratorTest.cpp

TEST(ZipIteratorTest, Basic) {
   for (auto [idx, elem] : enumerate(zip(odd, shortArr))) {
     EXPECT_LT(idx, static_cast<size_t>(2));
   }
   for (auto [idx, elem] : enumerate(zip(shortArr, odd))) {
     EXPECT_LT(idx, static_cast<size_t>(2));
   }
 }
 
+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};
   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;