Add `all_of_zip` to STLExtras

[mehdi_amini]

Summary

This takes two ranges and invokes a predicate on the element-wise pair in the
ranges. It returns true if all the pairs are matching the predicate and the ranges
have the same size.
It is useful with containers that aren't random iterator where we can't check the
sizes in O(1).

[rriddle]

llvm/include/llvm/ADT/STLExtras.h

1724

The name here feels a bit weird. This is essentially just std::ranges::equal/std::equal with special size mismatch behavior. Can we name it something closer to what it would be replaced by? Also, in what situations is the size mismatch result needed (i.e. the llvm::None)?

[mehdi_amini]

llvm/include/llvm/ADT/STLExtras.h

1724

I was trying to differentiate between "length mismatch" and "element mismatch", but maybe we don't need to?

[dblaikie]

Seems like a bit of a narrow use case combination of two features - is there much missing/a way we could generalize the more general tools so they could be composed easily enough in the cases where the combination is desirable?

[mehdi_amini]

The issue is that I think the "expected" way in C++ to have a zip operator that fails when the length mismatch is to use exceptions. I'm not how else to have a range that can fail when iterated on?

[dblaikie]

In D106605#2905907, @mehdi_amini wrote:

The issue is that I think the "expected" way in C++ to have a zip operator that fails when the length mismatch is to use exceptions. I'm not how else to have a range that can fail when iterated on?

Yeah, while keeping the iterator interface (rather than using a different interface, which would then be less composable) and with arbitrary elements (so you can't construct a default/failed object) it's a bit limited.

What about if the range itself was boolean testable (or had some "areBothRangesTraversed" operation or the like) that could be queried (assert-fails if queried early) after the zip has reached the end (and in that case the end would be the end of either range, whichever comes first)?

[mehdi_amini]

What about if the range itself was boolean testable (or had some "areBothRangesTraversed" operation or the like) that could be queried (assert-fails if queried early) after the zip has reached the end (and in that case the end would be the end of either range, whichever comes first)?

That's interesting: the range itself is disconnected from the iterator state though, so I don't think you could just query the range.

The best I can think of would be something that would end-up looking like:

auto range = zip(a, b, c);
auto it = range.begin();
for (; it != range.end(); ++it) {
   ...
}
if (!it.all_equals(range.end())  // bike shed the API name
   // error

[mehdi_amini]

Augment zip_common with a comparator to test for complete equality of the zipped iterators

[rriddle]

Would it be possible to use this to remove this hacky thing I did a while ago?

https://github.com/llvm/llvm-project/blob/fac6c5cf9b65fba6e028d9d2fdbec5c3114db542/mlir/include/mlir/IR/OpBase.td#L2371

[mehdi_amini]

I suspect for this case you just want this D106913 right?

[rriddle]

In D106605#2908386, @mehdi_amini wrote:

I suspect for this case you just want this D106913 right?

LGTM

[dblaikie]

In D106605#2908016, @mehdi_amini wrote:

What about if the range itself was boolean testable (or had some "areBothRangesTraversed" operation or the like) that could be queried (assert-fails if queried early) after the zip has reached the end (and in that case the end would be the end of either range, whichever comes first)?

That's interesting: the range itself is disconnected from the iterator state though, so I don't think you could just query the range.

The best I can think of would be something that would end-up looking like:

auto range = zip(a, b, c);
auto it = range.begin();
for (; it != range.end(); ++it) {
   ...
}
if (!it.all_equals(range.end())  // bike shed the API name
   // error

Ah, yeah, fair enough. Looks OK.

Though I worry about the likely combinatorial expansion of the ultimate API, still - all_of_zip implies the likely existence of none_of_zip and any_of_zip - which admittedly can be defined in terms of one another, so it's not loads of duplication... eh, maybe that's OK then, if they get created.

Carry on!

llvm/include/llvm/ADT/STLExtras.h

decltype(iterators) tup_inc(std::index_sequence<Ns...>) const {
     return std::tuple<Iters...>(std::next(std::get<Ns>(iterators))...);
   }
 
   template <size_t... Ns>
   decltype(iterators) tup_dec(std::index_sequence<Ns...>) const {
     return std::tuple<Iters...>(std::prev(std::get<Ns>(iterators))...);
   }
 
+  template <size_t... Ns>
+  bool test_all_equals(const zip_common &other,
+            std::index_sequence<Ns...>) const {
+    return all_of(std::initializer_list<bool>{std::get<Ns>(this->iterators) ==
+                                              std::get<Ns>(other.iterators)...},
+                  identity<bool>{});
+  }
+
 public:
   zip_common(Iters &&... ts) : iterators(std::forward<Iters>(ts)...) {}
 
   value_type operator*() { return deref(std::index_sequence_for<Iters...>{}); }
 
   const value_type operator*() const {
     return deref(std::index_sequence_for<Iters...>{});
   }
 
   ZipType &operator++() {
     iterators = tup_inc(std::index_sequence_for<Iters...>{});
     return *reinterpret_cast<ZipType *>(this);
   }
 
   ZipType &operator--() {
     static_assert(Base::IsBidirectional,
                   "All inner iterators must be at least bidirectional.");
     iterators = tup_dec(std::index_sequence_for<Iters...>{});
     return *reinterpret_cast<ZipType *>(this);
   }
+
+  /// Return true if all the iterator are matching `other`'s iterators.
+  bool all_equals(zip_common &other) {
+    return test_all_equals(other, std::index_sequence_for<Iters...>{});
+  }
 };
 
 template <typename... Iters>
 struct zip_first : public zip_common<zip_first<Iters...>, Iters...> {
   using Base = zip_common<zip_first<Iters...>, Iters...>;
 
   bool operator==(const zip_first<Iters...> &other) const {
     return std::get<0>(this->iterators) == std::get<0>(other.iterators);
bool is_splat(R &&Range) {
   return range_size != 0 && (range_size == 1 ||
          std::equal(adl_begin(Range) + 1, adl_end(Range), adl_begin(Range)));
 }
 
 /// Provide a container algorithm similar to C++ Library Fundamentals v2's
 /// `erase_if` which is equivalent to:
 ///
 ///   C.erase(remove_if(C, pred), C.end());
 ///
 /// This version works for any container with an erase method call accepting
 /// two iterators.
 template <typename Container, typename UnaryPredicate>
 void erase_if(Container &C, UnaryPredicate P) {
   C.erase(remove_if(C, P), C.end());
 }
 
 /// Wrapper function to remove a value from a container:

 decltype(auto) apply_tuple(F &&f, Tuple &&t) {
   using Indices = std::make_index_sequence<
       std::tuple_size<typename std::decay<Tuple>::type>::value>;
 
   return detail::apply_tuple_impl(std::forward<F>(f), std::forward<Tuple>(t),
                                   Indices{});
 }
 
+namespace detail {
+
+template <typename Predicate, typename... Args>
+bool all_of_zip_predicate_first(Predicate &&P, Args &&...args) {
+  auto z = zip(args...);
+  auto it = z.begin();
+  auto end = z.end();
+  while (it != end) {
+    if (!apply_tuple([&](auto &&...args) { return P(args...); }, *it))
+      return false;
+    ++it;
+  }
+  return it.all_equals(end);
+}
+
+// Just an adaptor to switch the order of argument and have the predicate before
+// the zipped inputs.
+template <typename... ArgsThenPredicate, size_t... InputIndexes>
+bool all_of_zip_predicate_last(
+    std::tuple<ArgsThenPredicate...> argsThenPredicate,
+    std::index_sequence<InputIndexes...>) {
+  auto constexpr OutputIndex =
+      std::tuple_size<decltype(argsThenPredicate)>::value - 1;
+  return all_of_zip_predicate_first(std::get<OutputIndex>(argsThenPredicate),
+                             std::get<InputIndexes>(argsThenPredicate)...);
+}
+
+} // end namespace detail
+
+/// Compare two zipped ranges using the provided predicate (as last argument).
+/// Return true if all elements satisfy the predicate and false otherwise.
+//  Return false if the zipped iterator aren't all at end (size mismatch).
+template <typename... ArgsAndPredicate>
+bool all_of_zip(ArgsAndPredicate &&...argsAndPredicate) {
+  return detail::all_of_zip_predicate_last(
+      std::forward_as_tuple(argsAndPredicate...),
+      std::make_index_sequence<sizeof...(argsAndPredicate) - 1>{});
+}
+
 /// Return true if the sequence [Begin, End) has exactly N items. Runs in O(N)
 /// time. Not meant for use with random-access iterators.
 /// Can optionally take a predicate to filter lazily some items.
 template <typename IterTy,
           typename Pred = bool (*)(const decltype(*std::declval<IterTy>()) &)>
 bool hasNItems(
     IterTy &&Begin, IterTy &&End, unsigned N,
     Pred &&ShouldBeCounted =

llvm/unittests/ADT/STLExtrasTest.cpp

int Destructors = 0;
     EXPECT_EQ(0, Moves);
     EXPECT_EQ(1, Destructors);
   }
   EXPECT_EQ(1, Copies);
   EXPECT_EQ(0, Moves);
   EXPECT_EQ(2, Destructors);
 }
 
+TEST(STLExtrasTest, AllOfZip) {
+  std::vector<int> v1 = {0, 4, 2, 1};
+  std::vector<int> v2 = {1, 4, 3, 6};
+  EXPECT_TRUE(all_of_zip(v1, v2, [](int v1, int v2) { return v1 <= v2; }));
+  EXPECT_FALSE(all_of_zip(v1, v2, [](int L, int R) { return L < R; }));
+
+  // Triple vectors
+  std::vector<int> v3 = {1, 6, 5, 7};
+  EXPECT_EQ(true, all_of_zip(v1, v2, v3, [](int a, int b, int c) {
+              return a <= b && b <= c;
+            }));
+  EXPECT_EQ(false, all_of_zip(v1, v2, v3, [](int a, int b, int c) {
+              return a < b && b < c;
+            }));
+
+  // Shorter vector should fail even with an always-true predicate.
+  std::vector<int> v_short = {1, 4};
+  EXPECT_EQ(false, all_of_zip(v1, v_short, [](int, int) { return true; }));
+  EXPECT_EQ(false,
+            all_of_zip(v1, v2, v_short, [](int, int, int) { return true; }));
+}
+
 } // namespace