Index: llvm/include/llvm/ADT/SmallVector.h =================================================================== --- llvm/include/llvm/ADT/SmallVector.h +++ llvm/include/llvm/ADT/SmallVector.h @@ -222,8 +222,10 @@ /// Reserve enough space to add one element, and return the updated element /// pointer in case it was a reference to the storage. - template static T *reserveForAndGetAddressImpl(U *This, T &Elt) { - if (LLVM_LIKELY(This->size() < This->capacity())) + template + static T *reserveForAndGetAddressImpl(U *This, T &Elt, size_t N) { + size_t NewSize = This->size() + N; + if (LLVM_LIKELY(NewSize <= This->capacity())) return const_cast(&Elt); bool ReferencesStorage = false; @@ -232,7 +234,7 @@ ReferencesStorage = true; Index = &Elt - This->begin(); } - This->grow(); + This->grow(NewSize); return ReferencesStorage ? This->begin() + Index : const_cast(&Elt); } @@ -356,14 +358,14 @@ /// Reserve enough space to add one element, and return the updated element /// pointer in case it was a reference to the storage. - T *reserveForAndGetAddress(T &Elt) { - return this->reserveForAndGetAddressImpl(this, Elt); + T *reserveForAndGetAddress(T &Elt, size_t N = 1) { + return this->reserveForAndGetAddressImpl(this, Elt, N); } /// Reserve enough space to add one element, and return the updated element /// pointer in case it was a reference to the storage. - const T *reserveForAndGetAddress(const T &Elt) { - return this->reserveForAndGetAddress(const_cast(Elt)); + const T *reserveForAndGetAddress(const T &Elt, size_t N = 1) { + return this->reserveForAndGetAddress(const_cast(Elt), N); } public: @@ -479,14 +481,14 @@ /// Reserve enough space to add one element, and return the updated element /// pointer in case it was a reference to the storage. - T *reserveForAndGetAddress(T &Elt) { - return this->reserveForAndGetAddressImpl(this, Elt); + T *reserveForAndGetAddress(T &Elt, size_t N = 1) { + return this->reserveForAndGetAddressImpl(this, Elt, N); } /// Reserve enough space to add one element, and return the updated element /// pointer in case it was a reference to the storage. - const T *reserveForAndGetAddress(const T &Elt) { - return this->reserveForAndGetAddress(const_cast(Elt)); + const T *reserveForAndGetAddress(const T &Elt, size_t N = 1) { + return this->reserveForAndGetAddress(const_cast(Elt), N); } public: @@ -609,12 +611,9 @@ } /// Append \p NumInputs copies of \p Elt to the end. - void append(size_type NumInputs, const T &Elt) { - this->assertSafeToAdd(&Elt, NumInputs); - if (NumInputs > this->capacity() - this->size()) - this->grow(this->size()+NumInputs); - - std::uninitialized_fill_n(this->end(), NumInputs, Elt); + void append(size_type NumInputs, ValueParamT Elt) { + const T *EltPtr = this->reserveForAndGetAddress(Elt, NumInputs); + std::uninitialized_fill_n(this->end(), NumInputs, *EltPtr); this->set_size(this->size() + NumInputs); } @@ -741,7 +740,7 @@ return insert_one_maybe_copy(I, Elt); } - iterator insert(iterator I, size_type NumToInsert, const T &Elt) { + iterator insert(iterator I, size_type NumToInsert, ValueParamT Elt) { // Convert iterator to elt# to avoid invalidating iterator when we reserve() size_t InsertElt = I - this->begin(); @@ -752,11 +751,9 @@ assert(this->isReferenceToStorage(I) && "Insertion iterator is out of bounds."); - // Check that adding NumToInsert elements won't invalidate Elt. - this->assertSafeToAdd(&Elt, NumToInsert); - - // Ensure there is enough space. - reserve(this->size() + NumToInsert); + // Ensure there is enough space, and get the (maybe updated) address of + // Elt. + const T *EltPtr = this->reserveForAndGetAddress(Elt, NumToInsert); // Uninvalidate the iterator. I = this->begin()+InsertElt; @@ -773,7 +770,12 @@ // Copy the existing elements that get replaced. std::move_backward(I, OldEnd-NumToInsert, OldEnd); - std::fill_n(I, NumToInsert, Elt); + // If we just moved the element we're inserting, be sure to update + // the reference (never happens if TakesParamByValue). + if (!TakesParamByValue && I <= EltPtr && EltPtr < this->end()) + EltPtr += NumToInsert; + + std::fill_n(I, NumToInsert, *EltPtr); return I; } @@ -786,11 +788,16 @@ size_t NumOverwritten = OldEnd-I; this->uninitialized_move(I, OldEnd, this->end()-NumOverwritten); + // If we just moved the element we're inserting, be sure to update + // the reference (never happens if TakesParamByValue). + if (!TakesParamByValue && I <= EltPtr && EltPtr < this->end()) + EltPtr += NumToInsert; + // Replace the overwritten part. - std::fill_n(I, NumOverwritten, Elt); + std::fill_n(I, NumOverwritten, *EltPtr); // Insert the non-overwritten middle part. - std::uninitialized_fill_n(OldEnd, NumToInsert-NumOverwritten, Elt); + std::uninitialized_fill_n(OldEnd, NumToInsert - NumOverwritten, *EltPtr); return I; } Index: llvm/unittests/ADT/SmallVectorTest.cpp =================================================================== --- llvm/unittests/ADT/SmallVectorTest.cpp +++ llvm/unittests/ADT/SmallVectorTest.cpp @@ -1113,9 +1113,15 @@ TYPED_TEST(SmallVectorReferenceInvalidationTest, Append) { auto &V = this->V; (void)V; -#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST - EXPECT_DEATH(V.append(1, V.back()), this->AssertionMessage); -#endif + V.append(1, V.back()); + int N = this->NumBuiltinElts(V); + EXPECT_EQ(N, V[N - 1]); + + // Append enough more elements that V will grow again. If reference + // invalidation breaks in the future, sanitizers should be able to catch a + // use-after-free here. + V.append(V.capacity(), V.front()); + EXPECT_EQ(1, V.back()); } TYPED_TEST(SmallVectorReferenceInvalidationTest, AppendRange) { @@ -1198,9 +1204,17 @@ TYPED_TEST(SmallVectorReferenceInvalidationTest, InsertN) { auto &V = this->V; (void)V; -#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST - EXPECT_DEATH(V.insert(V.begin(), 2, V.back()), this->AssertionMessage); -#endif + + // Cover NumToInsert <= this->end() - I. + V.insert(V.begin() + 1, 1, V.back()); + int N = this->NumBuiltinElts(V); + EXPECT_EQ(N, V[1]); + + // Cover NumToInsert > this->end() - I. Also insert enough more elements that + // V will grow again. If reference invalidation breaks in the future, + // sanitizers should be able to catch a use-after-free here. + V.insert(V.begin(), V.capacity(), V.front()); + EXPECT_EQ(1, V.front()); } TYPED_TEST(SmallVectorReferenceInvalidationTest, InsertRange) {