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 @@ -720,12 +720,15 @@ EarlyIncIteratorT(std::end(std::forward(Range)))); } -// forward declarations required by zip_shortest/zip_first/zip_longest +// Forward declarations required by zip_shortest/zip_equal/zip_first/zip_longest template bool all_of(R &&range, UnaryPredicate P); + template bool any_of(R &&range, UnaryPredicate P); +template bool all_equal(std::initializer_list Values); + namespace detail { using std::declval; @@ -879,6 +882,20 @@ std::forward(t), std::forward(u), std::forward(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 +detail::zippy 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( + std::forward(t), std::forward(u), std::forward(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. diff --git a/llvm/unittests/ADT/IteratorTest.cpp b/llvm/unittests/ADT/IteratorTest.cpp --- a/llvm/unittests/ADT/IteratorTest.cpp +++ b/llvm/unittests/ADT/IteratorTest.cpp @@ -416,6 +416,33 @@ } } +TEST(ZipIteratorTest, ZipEqualBasic) { + const SmallVector pi = {3, 1, 4, 1, 5, 8}; + const SmallVector 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 pi = {3, 1, 4, 1, 5, 8}; + const SmallVector 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 pi{3, 1, 4, 1, 5, 9};