Index: llvm/include/llvm/ADT/STLExtras.h =================================================================== --- llvm/include/llvm/ADT/STLExtras.h +++ llvm/include/llvm/ADT/STLExtras.h @@ -418,6 +418,61 @@ std::end(std::forward(Range)), Pred)); } +/// An iterator adaptor that does an early increment of the iterator immediately +/// after it is dereferenced. +/// +template +class early_inc_iterator + : public iterator_adaptor_base, + WrappedIteratorT, + std::forward_iterator_tag> { + using BaseT = + iterator_adaptor_base, + WrappedIteratorT, std::forward_iterator_tag>; + + using PointerT = typename std::iterator_traits::pointer; + +protected: + WrappedIteratorT NextI; + +public: + // Construct the iterator. The begin iterator needs to know where the end + // is, so that it can properly stop when it gets there. The end iterator only + // needs the predicate to support bidirectional iteration. + early_inc_iterator(WrappedIteratorT I) : BaseT(I), NextI(I) {} + + using BaseT::operator*; + typename BaseT::reference operator*() { + // The first time we dereference the iterator, we know we're not at the end + // so we can compute the next iterator. We do this early (here), prior to + // any possible mutation of the sequence we are iterating over by the user + // after the reference to this node is returned. + if (this->I == NextI) + ++NextI; + + return BaseT::operator*(); + } + + using BaseT::operator++; + early_inc_iterator &operator++() { + // If we haven't yet computed the next iterator, do so. + if (this->I == NextI) + ++NextI; + + // And mark that this is now the current iterator. + this->I = NextI; + return *this; + } +}; + +template +iterator_range>> +make_early_inc_range(RangeT &&Range) { + using EarlyIncIteratorT = early_inc_iterator>; + return make_range(EarlyIncIteratorT(std::begin(std::forward(Range))), + EarlyIncIteratorT(std::end(std::forward(Range)))); +} + // forward declarations required by zip_shortest/zip_first template bool all_of(R &&range, UnaryPredicate P); Index: llvm/unittests/ADT/STLExtrasTest.cpp =================================================================== --- llvm/unittests/ADT/STLExtrasTest.cpp +++ llvm/unittests/ADT/STLExtrasTest.cpp @@ -364,4 +364,38 @@ EXPECT_EQ(5, count); } +TEST(STLExtrasTest, EarlyIncrementTest) { + std::list L = {1, 2, 3, 4}; + + auto EIR = make_early_inc_range(L); + + auto I = EIR.begin(); + auto EI = EIR.end(); + EXPECT_NE(I, EI); + + EXPECT_EQ(1, *I); + // Repeated dereferences work. + EXPECT_EQ(1, *I); + + // Basic increment works, including repeated dereferences. + ++I; + EXPECT_NE(I, EI); + EXPECT_EQ(2, *I); + EXPECT_EQ(2, *I); + + // Inserting shouldn't break anything. We should be able to keep dereferencing + // the currrent iterator and increment. The increment to go to the "next" + // iterator from before we inserted. + L.insert(std::next(L.begin(), 2), -1); + EXPECT_EQ(2, *I); + ++I; + EXPECT_EQ(3, *I); + EXPECT_EQ(3, *I); + + // We should also be able to increment past something without looking at it. + ++I; + ++I; + EXPECT_EQ(EIR.end(), I); +} + } // namespace