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,36 @@ EXPECT_EQ(iters, 4u); } +TEST(ZipIteratorTest, ZipLongestBasic) { + using namespace std; + const vector pi{3, 1, 4, 1, 5, 9}; + const vector e{2, 7, 1, 8}; + + { + // Check left range longer than right. + const vector, Optional>> expected{ + {3, 2}, {1, 7}, {4, 1}, {1, 8}, {5, None}, {9, None}}; + int iters = 0; + for (auto tup : zip_longest(pi, e)) { + EXPECT_EQ(tup, expected[iters]); + iters += 1; + } + EXPECT_EQ(iters, expected.size()); + } + + { + // Check right range longer than left. + const vector, Optional>> expected{ + {2, 3}, {7, 1}, {1, 4}, {8, 1}, {None, 5}, {None, 9}}; + int iters = 0; + for (auto tup : zip_longest(e, pi)) { + EXPECT_EQ(tup, expected[iters]); + iters += 1; + } + EXPECT_EQ(iters, expected.size()); + } +} + TEST(ZipIteratorTest, Mutability) { using namespace std; const SmallVector pi{3, 1, 4, 1, 5, 9};