diff --git a/llvm/include/llvm/ADT/STLExtras.h b/llvm/include/llvm/ADT/STLExtras.h --- a/llvm/include/llvm/ADT/STLExtras.h +++ b/llvm/include/llvm/ADT/STLExtras.h @@ -630,6 +630,14 @@ return std::tuple(std::prev(std::get(iterators))...); } + template + bool test_all_equals(const zip_common &other, + std::index_sequence) const { + return all_of(std::initializer_list{std::get(this->iterators) == + std::get(other.iterators)...}, + identity{}); + } + public: zip_common(Iters &&... ts) : iterators(std::forward(ts)...) {} @@ -650,6 +658,11 @@ iterators = tup_dec(std::index_sequence_for{}); return *reinterpret_cast(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{}); + } }; template @@ -1986,6 +1999,45 @@ Indices{}); } +namespace detail { + +template +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 +bool all_of_zip_predicate_last( + std::tuple argsThenPredicate, + std::index_sequence) { + auto constexpr OutputIndex = + std::tuple_size::value - 1; + return all_of_zip_predicate_first(std::get(argsThenPredicate), + std::get(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 +bool all_of_zip(ArgsAndPredicate &&...argsAndPredicate) { + return detail::all_of_zip_predicate_last( + std::forward_as_tuple(argsAndPredicate...), + std::make_index_sequence{}); +} + /// 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. diff --git a/llvm/unittests/ADT/STLExtrasTest.cpp b/llvm/unittests/ADT/STLExtrasTest.cpp --- a/llvm/unittests/ADT/STLExtrasTest.cpp +++ b/llvm/unittests/ADT/STLExtrasTest.cpp @@ -876,4 +876,26 @@ EXPECT_EQ(2, Destructors); } +TEST(STLExtrasTest, AllOfZip) { + std::vector v1 = {0, 4, 2, 1}; + std::vector 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 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 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