Index: include/llvm/ADT/STLExtras.h =================================================================== --- include/llvm/ADT/STLExtras.h +++ include/llvm/ADT/STLExtras.h @@ -508,9 +508,11 @@ EarlyIncIteratorT(std::end(std::forward(Range)))); } -// forward declarations required by zip_shortest/zip_first +// forward declarations required by zip_shortest/zip_first/zip_longest template bool all_of(R &&range, UnaryPredicate P); +template +bool any_of(R &&range, UnaryPredicate P); template struct index_sequence; @@ -661,6 +663,145 @@ std::forward(t), std::forward(u), std::forward(args)...); } +namespace detail { +template +static Iter next_or_end(const Iter &I, const Iter &End) { + if (I == End) + return End; + return std::next(I); +} + +template +static auto deref_or_none(const Iter &I, const Iter &End) + -> llvm::Optional::type>::type> { + if (I == End) + return None; + return *I; +} + +template struct ZipLongestOptionalValueType { + using type = std::tuple< + llvm::Optional())>::type>::type>...>; +}; + +template +class zip_longest_iterator + : public iterator_facade_base< + zip_longest_iterator, + typename std::common_type< + std::forward_iterator_tag, + typename std::iterator_traits::iterator_category...>::type, + typename ZipLongestOptionalValueType::type, + typename std::iterator_traits>::type>::difference_type, + typename ZipLongestOptionalValueType::type *, + typename ZipLongestOptionalValueType::type> { +public: + using value_type = typename ZipLongestOptionalValueType::type; + +private: + std::tuple iterators; + std::tuple end_iterators; + + template + bool test(const zip_longest_iterator &other, + index_sequence) const { + return llvm::any_of( + std::initializer_list{std::get(this->iterators) != + std::get(other.iterators)...}, + identity{}); + } + + template value_type deref(index_sequence) const { + return value_type( + deref_or_none(std::get(iterators), std::get(end_iterators))...); + } + + template + decltype(iterators) tup_inc(index_sequence) const { + return std::tuple( + next_or_end(std::get(iterators), std::get(end_iterators))...); + } + +public: + zip_longest_iterator(std::pair... ts) + : iterators(std::forward(ts.first)...), + end_iterators(std::forward(ts.second)...) {} + + value_type operator*() { return deref(index_sequence_for{}); } + + value_type operator*() const { return deref(index_sequence_for{}); } + + zip_longest_iterator &operator++() { + iterators = tup_inc(index_sequence_for{}); + return *this; + } + + bool operator==(const zip_longest_iterator &other) const { + return !test(other, index_sequence_for{}); + } +}; + +template class zip_longest_range { +public: + using iterator = + zip_longest_iterator()))...>; + using iterator_category = typename iterator::iterator_category; + using value_type = typename iterator::value_type; + using difference_type = typename iterator::difference_type; + using pointer = typename iterator::pointer; + using reference = typename iterator::reference; + +private: + std::tuple ts; + + template iterator begin_impl(index_sequence) const { + return iterator(std::make_pair(adl_begin(std::get(ts)), + adl_end(std::get(ts)))...); + } + + template iterator end_impl(index_sequence) const { + return iterator(std::make_pair(adl_end(std::get(ts)), + adl_end(std::get(ts)))...); + } + +public: + zip_longest_range(Args &&... ts_) : ts(std::forward(ts_)...) {} + + iterator begin() const { return begin_impl(index_sequence_for{}); } + iterator end() const { return end_impl(index_sequence_for{}); } +}; +} // namespace detail + +/// Iterate over two or more iterators at the same time. Iteration continues +/// until all iterators reach the end. The llvm::Optional only contains a value +/// if the iterator has not reached the end. +template +detail::zip_longest_range zip_longest(T &&t, U &&u, + Args &&... args) { + return detail::zip_longest_range( + std::forward(t), std::forward(u), std::forward(args)...); +} + +namespace detail { +template +static auto deref_or_default(const Iter &I, const Iter &End) -> + typename std::remove_const< + typename std::remove_reference::type>::type { + if (I == End) + return {}; // value-initialize + return *I; +} + +template struct ZipLongestDefaultValueType { + using type = + std::tuple())>::type>::type...>; +}; +} // namespace detail + /// Iterator wrapper that concatenates sequences together. /// /// This can concatenate different iterators, even with different types, into Index: unittests/ADT/IteratorTest.cpp =================================================================== --- unittests/ADT/IteratorTest.cpp +++ unittests/ADT/IteratorTest.cpp @@ -328,6 +328,50 @@ EXPECT_EQ(iters, 4u); } +TEST(ZipIteratorTest, ZipLongestBasic) { + using namespace std; + const SmallVector pi{3, 1, 4, 1, 5, 9}; + const SmallVector lsb{1, 1, 0, 1, 1, 1, 0, 0}; + + { + int total_iters = 0; + int left_iters = 0; + int right_iters = 0; + for (auto tup : zip_longest(pi, lsb)) { + total_iters += 1; + if (get<0>(tup).hasValue()) + left_iters += 1; + if (get<1>(tup).hasValue()) + right_iters += 1; + if (!get<0>(tup) || !get<1>(tup)) + continue; + EXPECT_EQ(get<0>(tup).getValue() & 0x01, get<1>(tup).getValue()); + } + EXPECT_EQ(total_iters, 8); + EXPECT_EQ(left_iters, 6); + EXPECT_EQ(right_iters, 8); + } + + { + int total_iters = 0; + int left_iters = 0; + int right_iters = 0; + for (auto tup : zip_longest(lsb, pi)) { + total_iters += 1; + if (get<0>(tup).hasValue()) + left_iters += 1; + if (get<1>(tup).hasValue()) + right_iters += 1; + if (!get<0>(tup) || !get<1>(tup)) + continue; + EXPECT_EQ(get<0>(tup).getValue(), get<1>(tup).getValue() & 0x01); + } + EXPECT_EQ(total_iters, 8); + EXPECT_EQ(left_iters, 8); + EXPECT_EQ(right_iters, 6); + } +} + TEST(ZipIteratorTest, Mutability) { using namespace std; const SmallVector pi{3, 1, 4, 1, 5, 9}; @@ -382,6 +426,30 @@ EXPECT_TRUE(all_of(pi, [](unsigned n) { return (n & 0x01) == 0; })); } +TEST(ZipIteratorTest, ZipLongestFilter) { + using namespace std; + vector pi{3, 1, 4, 1, 5, 9}; + vector e{2, 7, 1, 8}; + vector, Optional>> expected{ + {1, 7}, {1, 8}, {9, None}}; + + auto zipped = zip_longest(pi, e); + auto filtered = make_filter_range(zipped, [](decltype(zipped)::value_type t) { + unsigned digitsum = get<0>(t).getValueOr(0) + get<1>(t).getValueOr(0); + return digitsum != 5; + }); + + int iters = 0; + for (auto tup : filtered) { + EXPECT_EQ(tup, expected[iters]); + iters += 1; + } + + // Should have skipped pi[0]/e[0], pi[2]/e[2] and pi[4]/e[4] + // (because their sum is 5). + EXPECT_EQ(iters, 3); +} + TEST(ZipIteratorTest, Reverse) { using namespace std; vector ascending{0, 1, 2, 3, 4, 5};