ADT: Fix reference invalidation in N-element SmallVector::append and insert

[dexonsmith]

Summary

For small enough, trivially copyable T, take the parameter by-value in
SmallVector::append and SmallVector::insert. Otherwise, when
growing, update the argument appropriately.

Depends on https://reviews.llvm.org/D93779.
Split out from https://reviews.llvm.org/D91837.

[dblaikie]

llvm/unittests/ADT/SmallVectorTest.cpp

1156

Might be worth making this V.capacity() - V.size() + 1 to make it more precise? While V.capacity() will always be enough, I think (to me at least) that might make the test a bit harder to read, since capacity() may be significantly more than needed.

[dexonsmith]

Thanks for the review! Landed in 3043e5a5c33c4c871f4a1dfd621a8839f9a1f0b3.

llvm/unittests/ADT/SmallVectorTest.cpp

1156

Good point. I updated as suggested in the commit.

In the similar case for the InsertN test below, I decided to expand the comment and leave it as V.capacity() since being precise there made it harder to reason about (since there we care about NumToInsert being larger than something).

llvm/include/llvm/ADT/SmallVector.h

void assertSafeToEmplace(ArgType1 &Arg1, ArgTypes &... Args) {
     this->assertSafeToAdd(&Arg1);
     this->assertSafeToEmplace(Args...);
   }
   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 <class U>
-  static const T *reserveForAndGetAddressImpl(U *This, const T &Elt) {
-    if (LLVM_LIKELY(This->size() < This->capacity()))
+  static const T *reserveForAndGetAddressImpl(U *This, const T &Elt, size_t N) {
+    size_t NewSize = This->size() + N;
+    if (LLVM_LIKELY(NewSize <= 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();
+    This->grow(NewSize);
     return ReferencesStorage ? This->begin() + Index : &Elt;
   }
 
 public:
   using size_type = size_t;
   using difference_type = ptrdiff_t;
   using value_type = T;
   using iterator = T *;
protected:
 
   /// Grow the allocated memory (without initializing new elements), doubling
   /// the size of the allocated memory. Guarantees space for at least one more
   /// 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);
+  const T *reserveForAndGetAddress(const 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.
-  T *reserveForAndGetAddress(T &Elt) {
-    return const_cast<T *>(this->reserveForAndGetAddressImpl(this, Elt));
+  T *reserveForAndGetAddress(T &Elt, size_t N = 1) {
+    return const_cast<T *>(this->reserveForAndGetAddressImpl(this, Elt, N));
   }
 
   static T &&forward_value_param(T &&V) { return std::move(V); }
   static const T &forward_value_param(const T &V) { return V; }
 
 public:
   void push_back(const T &Elt) {
     const T *EltPtr = reserveForAndGetAddress(Elt);
protected:
   }
 
   /// Double the size of the allocated memory, guaranteeing space for at
   /// 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);
+  const T *reserveForAndGetAddress(const 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.
-  T *reserveForAndGetAddress(T &Elt) {
-    return const_cast<T *>(this->reserveForAndGetAddressImpl(this, Elt));
+  T *reserveForAndGetAddress(T &Elt, size_t N = 1) {
+    return const_cast<T *>(this->reserveForAndGetAddressImpl(this, Elt, N));
   }
 
   /// Copy \p V or return a reference, depending on \a ValueParamT.
   static ValueParamT forward_value_param(ValueParamT V) { return V; }
 
 public:
   void push_back(ValueParamT Elt) {
     const T *EltPtr = reserveForAndGetAddress(Elt);
void append(in_iter in_start, in_iter in_end) {
     if (NumInputs > this->capacity() - this->size())
       this->grow(this->size()+NumInputs);
 
     this->uninitialized_copy(in_start, in_end, this->end());
     this->set_size(this->size() + NumInputs);
   }
 
   /// 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);
   }
 
   void append(std::initializer_list<T> IL) {
     append(IL.begin(), IL.end());
   }
 
   // FIXME: Consider assigning over existing elements, rather than clearing &
public:
   iterator insert(iterator I, T &&Elt) {
     return insert_one_impl(I, this->forward_value_param(std::move(Elt)));
   }
 
   iterator insert(iterator I, const T &Elt) {
     return insert_one_impl(I, this->forward_value_param(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();
 
     if (I == this->end()) {  // Important special case for empty vector.
       append(NumToInsert, Elt);
       return this->begin()+InsertElt;
     }
 
     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;
 
     // If there are more elements between the insertion point and the end of the
     // range than there are being inserted, we can use a simple approach to
     // insertion.  Since we already reserved space, we know that this won't
     // reallocate the vector.
     if (size_t(this->end()-I) >= NumToInsert) {
       T *OldEnd = this->end();
       append(std::move_iterator<iterator>(this->end() - NumToInsert),
              std::move_iterator<iterator>(this->end()));
 
       // 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;
     }
 
     // Otherwise, we're inserting more elements than exist already, and we're
     // not inserting at the end.
 
     // Move over the elements that we're about to overwrite.
     T *OldEnd = this->end();
     this->set_size(this->size() + NumToInsert);
     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;
   }
 
   template <typename ItTy,
             typename = std::enable_if_t<std::is_convertible<
                 typename std::iterator_traits<ItTy>::iterator_category,
                 std::input_iterator_tag>::value>>
   iterator insert(iterator I, ItTy From, ItTy To) {

llvm/unittests/ADT/SmallVectorTest.cpp

#endif
 
   // No assertion when shrinking, since the parameter isn't accessed.
   V.resize(N - 1, V.back());
 }
 
 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. This tests growing
+  // when already in large mode.
+  //
+  // If reference invalidation breaks in the future, sanitizers should be able
+  // to catch a use-after-free here.
+  V.append(V.capacity() - V.size() + 1, V.front());
+  EXPECT_EQ(1, V.back());
 }
 
 TYPED_TEST(SmallVectorReferenceInvalidationTest, AppendRange) {
   auto &V = this->V;
   (void)V;
 #if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
   EXPECT_DEATH(V.append(V.begin(), V.begin() + 1), this->AssertionMessage);
 
if (this->template isValueType<Constructable>()) {
     // Check the value got moved out.
     EXPECT_EQ(0, V.back());
   }
 }
 
 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, inserting enough elements that V will
+  // also grow again; V.capacity() will be more elements than necessary but
+  // it's a simple way to cover both conditions.
+  //
+  // 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) {
   auto &V = this->V;
   (void)V;
 #if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
   EXPECT_DEATH(V.insert(V.begin(), V.begin(), V.begin() + 1),
                this->AssertionMessage);