ADT: Fix reference invalidation in SmallVector::resize

[dexonsmith]

Summary

For small enough, trivially copyable T, take the parameter by-value in
SmallVector::resize. Otherwise, when growing, update the arugment
appropriately.

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

[dblaikie]

Looks good (did you happen to try this with sanitizers without the fix applied & observe a sanitizer failure when running the updated test? Would be nice to know that's actually working as intended)

[dexonsmith]

In D93781#2470743, @dblaikie wrote:

Looks good (did you happen to try this with sanitizers without the fix applied & observe a sanitizer failure when running the updated test? Would be nice to know that's actually working as intended)

No, but just now I ran with the SmallVector.h part of this patch mostly reverted (didn't add back the assertion), and the test fails:

/Users/dexonsmith/data/github/llvm-project/working/llvm/unittests/ADT/SmallVectorTest.cpp:1139: Failure
      Expected: N
      Which is: 3
To be equal to: V.back()
      Which is: 8-byte object <01-00 00-00 00-00 00-00>
/Users/dexonsmith/data/github/llvm-project/working/llvm/unittests/ADT/SmallVectorTest.cpp:1145: Failure
      Expected: 1
To be equal to: V.back()
      Which is: 8-byte object <01-04 00-00 00-00 00-00>

The line numbers are different vs. the patch posted here because of a rebase, but this seems to confirm the test is working even without sanitizers.

About to commit (check-llvm just finished on the rebase); I can follow up with ASan post-commit if you'd like (let me know if you would).

[dblaikie]

In D93781#2497219, @dexonsmith wrote:

In D93781#2470743, @dblaikie wrote:

Looks good (did you happen to try this with sanitizers without the fix applied & observe a sanitizer failure when running the updated test? Would be nice to know that's actually working as intended)

No, but just now I ran with the SmallVector.h part of this patch mostly reverted (didn't add back the assertion), and the test fails:

/Users/dexonsmith/data/github/llvm-project/working/llvm/unittests/ADT/SmallVectorTest.cpp:1139: Failure
      Expected: N
      Which is: 3
To be equal to: V.back()
      Which is: 8-byte object <01-00 00-00 00-00 00-00>
/Users/dexonsmith/data/github/llvm-project/working/llvm/unittests/ADT/SmallVectorTest.cpp:1145: Failure
      Expected: 1
To be equal to: V.back()
      Which is: 8-byte object <01-04 00-00 00-00 00-00>

The line numbers are different vs. the patch posted here because of a rebase, but this seems to confirm the test is working even without sanitizers.

About to commit (check-llvm just finished on the rebase); I can follow up with ASan post-commit if you'd like (let me know if you would).

Looks good to me as-is - thanks!

llvm/include/llvm/ADT/SmallVector.h

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.
   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, size_t N = 1) {
     return this->reserveForAndGetAddress(const_cast<T &>(Elt), N);
   }
 
 public:
private:
   }
 
 public:
   void resize(size_type N) { resizeImpl<false>(N); }
 
   /// Like resize, but \ref T is POD, the new values won't be initialized.
   void resize_for_overwrite(size_type N) { resizeImpl<true>(N); }
 
-  void resize(size_type N, const T &NV) {
+  void resize(size_type N, ValueParamT NV) {
     if (N == this->size())
       return;
 
     if (N < this->size()) {
       this->destroy_range(this->begin()+N, this->end());
       this->set_size(N);
       return;
     }
 
-    this->assertSafeToReferenceAfterResize(&NV, N);
-    if (this->capacity() < N)
-      this->grow(N);
-    std::uninitialized_fill(this->end(), this->begin() + N, NV);
-    this->set_size(N);
+    // N > this->size(). Defer to append.
+    this->append(N - this->size(), NV);
   }
 
   void reserve(size_type N) {
     if (this->capacity() < N)
       this->grow(N);
   }
 
   void pop_back_n(size_type NumItems) {

llvm/unittests/ADT/SmallVectorTest.cpp

EmplaceableArg<3> A3(true);
     EXPECT_TRUE(&back == &V.back());
     EXPECT_TRUE(V.size() == 1);
     EXPECT_TRUE(back.State == ES_Emplaced);
     EXPECT_TRUE(back.A0.State == EAS_Defaulted);
     EXPECT_TRUE(back.A1.State == EAS_Defaulted);
     EXPECT_TRUE(back.A2.State == EAS_Defaulted);
     EXPECT_TRUE(back.A3.State == EAS_Defaulted);
   }
   {
     SmallVector<Emplaceable, 3> V;
     Emplaceable &back = V.emplace_back(std::move(A0));
     EXPECT_TRUE(&back == &V.back());
     EXPECT_TRUE(V.size() == 1);
     EXPECT_TRUE(back.State == ES_Emplaced);
     EXPECT_TRUE(back.A0.State == EAS_RValue);
     EXPECT_TRUE(back.A1.State == EAS_Defaulted);
     EXPECT_TRUE(back.A2.State == EAS_Defaulted);
     EXPECT_TRUE(back.A3.State == EAS_Defaulted);
   }
   {
     SmallVector<Emplaceable, 3> V;
     Emplaceable &back = V.emplace_back(A0);
     EXPECT_TRUE(&back == &V.back());
     EXPECT_TRUE(V.size() == 1);
     EXPECT_TRUE(back.State == ES_Emplaced);
     EXPECT_TRUE(back.A0.State == EAS_LValue);
     EXPECT_TRUE(back.A1.State == EAS_Defaulted);
     EXPECT_TRUE(back.A2.State == EAS_Defaulted);
     EXPECT_TRUE(back.A3.State == EAS_Defaulted);
   }
   {
     SmallVector<Emplaceable, 3> V;
     Emplaceable &back = V.emplace_back(A0, A1);
     EXPECT_TRUE(&back == &V.back());
     EXPECT_TRUE(V.size() == 1);
     EXPECT_TRUE(back.State == ES_Emplaced);
     EXPECT_TRUE(back.A0.State == EAS_LValue);
     EXPECT_TRUE(back.A1.State == EAS_LValue);
     EXPECT_TRUE(back.A2.State == EAS_Defaulted);
TYPED_TEST(SmallVectorReferenceInvalidationTest, PushBackMoved) {
   if (this->template isValueType<Constructable>())
     EXPECT_EQ(0, V[N - 1]);
 }
 
 TYPED_TEST(SmallVectorReferenceInvalidationTest, Resize) {
   auto &V = this->V;
   (void)V;
   int N = this->NumBuiltinElts(V);
-#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
-  EXPECT_DEATH(V.resize(N + 1, V.back()), this->AssertionMessage);
-#endif
+  V.resize(N + 1, V.back());
+  EXPECT_EQ(N, V.back());
 
-  // No assertion when shrinking, since the parameter isn't accessed.
-  V.resize(N - 1, V.back());
+  // Resize to add enough 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.resize(V.capacity() + 1, V.front());
+  EXPECT_EQ(1, V.back());
 }
 
 TYPED_TEST(SmallVectorReferenceInvalidationTest, Append) {
   auto &V = this->V;
   (void)V;
   V.append(1, V.back());
   int N = this->NumBuiltinElts(V);
   EXPECT_EQ(N, V[N - 1]);