diff --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h --- a/llvm/include/llvm/ADT/SmallVector.h +++ b/llvm/include/llvm/ADT/SmallVector.h @@ -220,6 +220,23 @@ } void assertSafeToEmplace() {} + /// Reserve enough space to add one element, and return the updated element + /// pointer in case it was a reference to the storage. + template + static const T *reserveForAndGetAddressImpl(U *This, const T &Elt) { + if (LLVM_LIKELY(This->size() < This->capacity())) + return &Elt; + + bool ReferencesStorage = false; + int64_t Index = -1; + if (LLVM_UNLIKELY(This->isReferenceToStorage(&Elt))) { + ReferencesStorage = true; + Index = &Elt - This->begin(); + } + This->grow(); + return ReferencesStorage ? This->begin() + Index : &Elt; + } + public: using size_type = size_t; using difference_type = ptrdiff_t; @@ -303,7 +320,12 @@ (is_trivially_move_constructible::value) && std::is_trivially_destructible::value> class SmallVectorTemplateBase : public SmallVectorTemplateCommon { + friend class SmallVectorTemplateCommon; + protected: + static constexpr bool TakesParamByValue = false; + using ValueParamT = const T &; + SmallVectorTemplateBase(size_t Size) : SmallVectorTemplateCommon(Size) {} static void destroy_range(T *S, T *E) { @@ -333,20 +355,28 @@ /// element, or MinSize more elements if specified. void grow(size_t MinSize = 0); + /// 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->reserveForAndGetAddressImpl(this, Elt); + } + + /// 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 const_cast(this->reserveForAndGetAddressImpl(this, Elt)); + } + public: void push_back(const T &Elt) { - this->assertSafeToAdd(&Elt); - if (LLVM_UNLIKELY(this->size() >= this->capacity())) - this->grow(); - ::new ((void*) this->end()) T(Elt); + const T *EltPtr = reserveForAndGetAddress(Elt); + ::new ((void *)this->end()) T(*EltPtr); this->set_size(this->size() + 1); } void push_back(T &&Elt) { - this->assertSafeToAdd(&Elt); - if (LLVM_UNLIKELY(this->size() >= this->capacity())) - this->grow(); - ::new ((void*) this->end()) T(::std::move(Elt)); + T *EltPtr = reserveForAndGetAddress(Elt); + ::new ((void *)this->end()) T(::std::move(*EltPtr)); this->set_size(this->size() + 1); } @@ -396,7 +426,18 @@ /// skipping destruction. template class SmallVectorTemplateBase : public SmallVectorTemplateCommon { + friend class SmallVectorTemplateCommon; + protected: + /// True if it's cheap enough to take parameters by value. Doing so avoids + /// overhead related to mitigations for reference invalidation. + static constexpr bool TakesParamByValue = sizeof(T) <= 2 * sizeof(void *); + + /// Either const T& or T, depending on whether it's cheap enough to take + /// parameters by value. + using ValueParamT = + typename std::conditional::type; + SmallVectorTemplateBase(size_t Size) : SmallVectorTemplateCommon(Size) {} // No need to do a destroy loop for POD's. @@ -437,12 +478,22 @@ /// least one more element or MinSize if specified. void grow(size_t MinSize = 0) { this->grow_pod(MinSize, sizeof(T)); } + /// 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->reserveForAndGetAddressImpl(this, Elt); + } + + /// 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 const_cast(this->reserveForAndGetAddressImpl(this, Elt)); + } + public: - void push_back(const T &Elt) { - this->assertSafeToAdd(&Elt); - if (LLVM_UNLIKELY(this->size() >= this->capacity())) - this->grow(); - memcpy(reinterpret_cast(this->end()), &Elt, sizeof(T)); + void push_back(ValueParamT Elt) { + const T *EltPtr = reserveForAndGetAddress(Elt); + memcpy(reinterpret_cast(this->end()), EltPtr, sizeof(T)); this->set_size(this->size() + 1); } @@ -462,6 +513,9 @@ using size_type = typename SuperClass::size_type; protected: + using SmallVectorTemplateBase::TakesParamByValue; + using ValueParamT = typename SuperClass::ValueParamT; + // Default ctor - Initialize to empty. explicit SmallVectorImpl(unsigned N) : SmallVectorTemplateBase(N) {} @@ -628,6 +682,12 @@ private: template iterator insert_one_impl(iterator I, ArgType &&Elt) { + // Callers ensure that ArgType is derived from T. + static_assert( + std::is_same>, + T>::value, + "ArgType must be derived from T!"); + if (I == this->end()) { // Important special case for empty vector. this->push_back(::std::forward(Elt)); return this->end()-1; @@ -635,14 +695,11 @@ assert(this->isReferenceToStorage(I) && "Insertion iterator is out of bounds."); - // Check that adding an element won't invalidate Elt. - this->assertSafeToAdd(&Elt); - - if (this->size() >= this->capacity()) { - size_t EltNo = I-this->begin(); - this->grow(); - I = this->begin()+EltNo; - } + // Grow if necessary. + size_t Index = I - this->begin(); + std::remove_reference_t *EltPtr = + this->reserveForAndGetAddress(Elt); + I = this->begin() + Index; ::new ((void*) this->end()) T(::std::move(this->back())); // Push everything else over. @@ -650,21 +707,48 @@ this->set_size(this->size() + 1); // If we just moved the element we're inserting, be sure to update - // the reference. - std::remove_reference_t *EltPtr = &Elt; - if (this->isReferenceToRange(EltPtr, I, this->end())) + // the reference (never happens if TakesParamByValue). + static_assert(!TakesParamByValue || std::is_same::value, + "ArgType must be 'T' when taking by value!"); + if (!TakesParamByValue && this->isReferenceToRange(EltPtr, I, this->end())) ++EltPtr; *I = ::std::forward(*EltPtr); return I; } + template < + class ArgType, + std::enable_if_t< + std::is_same>, + T>::value && + !TakesParamByValue, + bool> = false> + iterator insert_one_maybe_copy(iterator I, ArgType &&Elt) { + return insert_one_impl(I, std::forward(Elt)); + } + + template < + class ArgType, + std::enable_if_t< + std::is_same>, + T>::value && + TakesParamByValue, + bool> = false> + iterator insert_one_maybe_copy(iterator I, ArgType &&Elt) { + // Copy Elt in order to mitigate reference invalidation without needing to + // update the pointer values in insert_one_impl. + return insert_one_impl(I, T(Elt)); + } + public: iterator insert(iterator I, T &&Elt) { - return insert_one_impl(I, std::move(Elt)); + return insert_one_maybe_copy(I, std::move(Elt)); } - iterator insert(iterator I, const T &Elt) { return insert_one_impl(I, Elt); } + iterator insert(iterator I, const T &Elt) { + return insert_one_maybe_copy(I, Elt); + } iterator insert(iterator I, size_type NumToInsert, const T &Elt) { // Convert iterator to elt# to avoid invalidating iterator when we reserve() diff --git a/llvm/unittests/ADT/SmallVectorTest.cpp b/llvm/unittests/ADT/SmallVectorTest.cpp --- a/llvm/unittests/ADT/SmallVectorTest.cpp +++ b/llvm/unittests/ADT/SmallVectorTest.cpp @@ -53,6 +53,7 @@ Constructable(Constructable && src) : constructed(true) { value = src.value; + src.value = 0; ++numConstructorCalls; ++numMoveConstructorCalls; } @@ -74,6 +75,7 @@ Constructable & operator=(Constructable && src) { EXPECT_TRUE(constructed); value = src.value; + src.value = 0; ++numAssignmentCalls; ++numMoveAssignmentCalls; return *this; @@ -1056,11 +1058,16 @@ return N; } + template static bool isValueType() { + return std::is_same::value; + } + void SetUp() override { SmallVectorTestBase::SetUp(); // Fill up the small size so that insertions move the elements. - V.append({0, 0, 0}); + for (int I = 0, E = NumBuiltinElts(V); I != E; ++I) + V.emplace_back(I + 1); } }; @@ -1074,19 +1081,54 @@ SmallVectorReferenceInvalidationTestTypes); TYPED_TEST(SmallVectorReferenceInvalidationTest, PushBack) { + // Note: setup adds [1, 2, ...] to V until it's at capacity in small mode. auto &V = this->V; - (void)V; -#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST - EXPECT_DEATH(V.push_back(V.back()), this->AssertionMessage); -#endif + int N = this->NumBuiltinElts(V); + + // Push back a reference to last element when growing from small storage. + V.push_back(V.back()); + EXPECT_EQ(N, V.back()); + + // Check that the old value is still there (not moved away). + EXPECT_EQ(N, V[V.size() - 2]); + + // Fill storage again. + V.back() = V.size(); + while (V.size() < V.capacity()) + V.push_back(V.size() + 1); + + // Push back a reference to last element when growing from large storage. + V.push_back(V.back()); + EXPECT_EQ(int(V.size()) - 1, V.back()); } TYPED_TEST(SmallVectorReferenceInvalidationTest, PushBackMoved) { + // Note: setup adds [1, 2, ...] to V until it's at capacity in small mode. auto &V = this->V; - (void)V; -#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST - EXPECT_DEATH(V.push_back(std::move(V.back())), this->AssertionMessage); -#endif + int N = this->NumBuiltinElts(V); + + // Push back a reference to last element when growing from small storage. + V.push_back(std::move(V.back())); + EXPECT_EQ(N, V.back()); + if (this->template isValueType()) { + // Check that the value was moved (not copied). + EXPECT_EQ(0, V[V.size() - 2]); + } + + // Fill storage again. + V.back() = V.size(); + while (V.size() < V.capacity()) + V.push_back(V.size() + 1); + + // Push back a reference to last element when growing from large storage. + V.push_back(std::move(V.back())); + + // Check the values. + EXPECT_EQ(int(V.size()) - 1, V.back()); + if (this->template isValueType()) { + // Check the value got moved out. + EXPECT_EQ(0, V[V.size() - 2]); + } } TYPED_TEST(SmallVectorReferenceInvalidationTest, Resize) { @@ -1150,20 +1192,53 @@ } TYPED_TEST(SmallVectorReferenceInvalidationTest, Insert) { + // Note: setup adds [1, 2, ...] to V until it's at capacity in small mode. auto &V = this->V; (void)V; -#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST - EXPECT_DEATH(V.insert(V.begin(), V.back()), this->AssertionMessage); -#endif + + // Insert a reference to the back (not at end() or else insert delegates to + // push_back()), growing out of small mode. Confirm the value was copied out + // (moving out Constructable sets it to 0). + V.insert(V.begin(), V.back()); + EXPECT_EQ(int(V.size() - 1), V.front()); + EXPECT_EQ(int(V.size() - 1), V.back()); + + // Fill up the vector again. + while (V.size() < V.capacity()) + V.push_back(V.size() + 1); + + // Grow again from large storage to large storage. + V.insert(V.begin(), V.back()); + EXPECT_EQ(int(V.size() - 1), V.front()); + EXPECT_EQ(int(V.size() - 1), V.back()); } TYPED_TEST(SmallVectorReferenceInvalidationTest, InsertMoved) { + // Note: setup adds [1, 2, ...] to V until it's at capacity in small mode. auto &V = this->V; (void)V; -#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST - EXPECT_DEATH(V.insert(V.begin(), std::move(V.back())), - this->AssertionMessage); -#endif + + // Insert a reference to the back (not at end() or else insert delegates to + // push_back()), growing out of small mode. Confirm the value was copied out + // (moving out Constructable sets it to 0). + V.insert(V.begin(), std::move(V.back())); + EXPECT_EQ(int(V.size() - 1), V.front()); + if (this->template isValueType()) { + // Check the value got moved out. + EXPECT_EQ(0, V.back()); + } + + // Fill up the vector again. + while (V.size() < V.capacity()) + V.push_back(V.size() + 1); + + // Grow again from large storage to large storage. + V.insert(V.begin(), std::move(V.back())); + EXPECT_EQ(int(V.size() - 1), V.front()); + if (this->template isValueType()) { + // Check the value got moved out. + EXPECT_EQ(0, V.back()); + } } TYPED_TEST(SmallVectorReferenceInvalidationTest, InsertN) {