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,243 @@ 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_optional_iterator + : public iterator_facade_base< + zip_longest_optional_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_optional_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_optional_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_optional_iterator &operator++() { + iterators = tup_inc(index_sequence_for{}); + return *this; + } + + bool operator==(const zip_longest_optional_iterator &other) const { + return !test(other, index_sequence_for{}); + } +}; + +template class zip_longest_optional_range { +public: + using iterator = zip_longest_optional_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(std::begin(std::get(ts)), + std::end(std::get(ts)))...); + } + + template iterator end_impl(index_sequence) const { + return iterator(std::make_pair(std::end(std::get(ts)), + std::end(std::get(ts)))...); + } + +public: + zip_longest_optional_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_optional_range +zip_longest_optional(T &&t, U &&u, Args &&... args) { + return detail::zip_longest_optional_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...>; +}; + +template +class zip_longest_default_iterator + : public iterator_facade_base< + zip_longest_default_iterator, + typename std::common_type< + std::forward_iterator_tag, + typename std::iterator_traits::iterator_category...>::type, + typename ZipLongestDefaultValueType::type, + typename std::iterator_traits>::type>::difference_type, + typename ZipLongestDefaultValueType::type *, + typename ZipLongestDefaultValueType::type> { +public: + using value_type = typename ZipLongestDefaultValueType::type; + +private: + std::tuple iterators; + std::tuple end_iterators; + + template + bool test(const zip_longest_default_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_default(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_default_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_default_iterator &operator++() { + iterators = tup_inc(index_sequence_for{}); + return *this; + } + + bool operator==(const zip_longest_default_iterator &other) const { + return !test(other, index_sequence_for{}); + } +}; + +template class zip_longest_default_range { +public: + using iterator = zip_longest_default_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(std::begin(std::get(ts)), + std::end(std::get(ts)))...); + } + + template iterator end_impl(index_sequence) const { + return iterator(std::make_pair(std::end(std::get(ts)), + std::end(std::get(ts)))...); + } + +public: + zip_longest_default_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 type's value-initialized default is +/// used for iterators that have reached the end when others did not. +template +detail::zip_longest_default_range +zip_longest_default(T &&t, U &&u, Args &&... args) { + return detail::zip_longest_default_range( + std::forward(t), std::forward(u), std::forward(args)...); +} + /// 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,84 @@ EXPECT_EQ(iters, 4u); } +TEST(ZipIteratorTest, ZipLongestOptionalBasic) { + 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_optional(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_optional(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, ZipLongestDefaultBasic) { + using namespace std; + const SmallVector pi{3, 1, 4, 1, 5, 9}; + const SmallVector ptr{&pi[0], &pi[1], &pi[2], &pi[3]}; + + { + int total_iters = 0; + int nondefault_iters = 0; + for (auto tup : zip_longest_default(pi, ptr)) { + total_iters += 1; + if (!get<1>(tup)) + continue; + nondefault_iters += 1; + EXPECT_EQ(get<0>(tup), *get<1>(tup)); + } + EXPECT_EQ(total_iters, 6); + EXPECT_EQ(nondefault_iters, 4); + } + + { + int total_iters = 0; + int nondefault_iters = 0; + for (auto tup : zip_longest_default(ptr, pi)) { + total_iters += 1; + if (!get<0>(tup)) + continue; + nondefault_iters += 1; + EXPECT_EQ(*get<0>(tup), get<1>(tup)); + } + EXPECT_EQ(total_iters, 6); + EXPECT_EQ(nondefault_iters, 4); + } +} + TEST(ZipIteratorTest, Mutability) { using namespace std; const SmallVector pi{3, 1, 4, 1, 5, 9}; @@ -382,6 +460,40 @@ EXPECT_TRUE(all_of(pi, [](unsigned n) { return (n & 0x01) == 0; })); } +TEST(ZipIteratorTest, ZipLongestOptionalFilter) { + using namespace std; + vector pi{3, 1, 4, 1, 5, 9}; + + int iters = 0; + auto zipped = zip_longest_optional(pi, vector{1, 1, 0, 1}); + for (auto tup : make_filter_range(zipped, [](decltype(zipped)::value_type t) { + return !get<1>(t).hasValue() || get<1>(t).getValue(); + })) { + iters += 1; + if (!get<0>(tup).hasValue() || !get<1>(tup).hasValue()) + continue; + EXPECT_EQ(get<0>(tup).getValue() & 0x01, get<1>(tup).getValue()); + } + + // Should have skipped pi[2]. + EXPECT_EQ(iters, 5); +} + +TEST(ZipIteratorTest, ZipLongestDefaultFilter) { + using namespace std; + vector pi{3, 1, 4, 1, 5, 9}; + + int iters = 0; + auto zipped = zip_longest_default(pi, vector{0, 0, 1, 0}); + for (auto tup : make_filter_range( + zipped, [](decltype(zipped)::value_type t) { return !get<1>(t); })) { + iters += 1; + EXPECT_EQ(~get<0>(tup) & 0x01, get<1>(tup)); + } + + EXPECT_EQ(iters, 5); +} + TEST(ZipIteratorTest, Reverse) { using namespace std; vector ascending{0, 1, 2, 3, 4, 5};