ADT: Fix reference invalidation in SmallVector::emplace_back and assign(N,V)

This fixes the final (I think?) reference invalidation in `SmallVector`
that we need to fix to align with `std::vector`. (There is still some
left in the range insert / append / assign, but the standard calls that
UB for `std::vector` so I think we don't care?)

For POD-like types, reimplement `emplace_back()` in terms of
`push_back()`, taking a copy even for large `T` rather than lose the
realloc optimization in `grow_pod()`.

For other types, split the grow operation in three and construct the new
element in the middle.

- `mallocForGrow()` calculates the new capacity and returns the result
  of `safe_malloc()`. We only need a single definition per
  `SmallVectorBase` so this is defined in SmallVector.cpp to avoid code
  size bloat. Moving this part of non-POD grow to the source file also
  allows the logic to be easily shared with `grow_pod`, and
  `report_size_overflow()` and `report_at_maximum_capacity()` can move
  there too.
- `moveElementsForGrow()` moves elements from the old to the new
  allocation.
- `takeAllocationForGrow()` frees the old allocation and saves the
  new allocation and capacity .

`SmallVector:assign(size_type, const T&)` also uses the split-grow
operations for non-POD, but it also has a semantic change when not
growing. Previously, assign would start with `clear()`, and so the old
elements were destructed and all elements of the new vector were
copy-constructed (potentially invalidating references). The new
implementation skips destruction and uses copy-assignment for the prefix
of the new vector that fits. The new semantics match what libc++ does
for `std::vector::assign()`.

Note that the following is another possible implementation:
```
  void assign(size_type NumElts, ValueParamT Elt) {
    std::fill_n(this->begin(), std::min(NumElts, this->size()), Elt);
    this->resize(NumElts, Elt);
  }
```
The downside of this simpler implementation is that if the vector has to
grow there will be `size()` redundant copy operations.

(I had planned on splitting this patch up into three for committing
(after getting performance numbers / initial review), but I've realized
that if this does for some reason need to be reverted we'll probably
want to revert the whole package...)

Differential Revision: https://reviews.llvm.org/D94739

1 files changed, 76 insertions, 20 deletions

diff --git a/llvm/unittests/ADT/SmallVectorTest.cpp b/llvm/unittests/ADT/SmallVectorTest.cpp
index 1f6c7db..b2cccc1 100644
--- a/llvm/unittests/ADT/SmallVectorTest.cpp
+++ b/llvm/unittests/ADT/SmallVectorTest.cpp
@@ -1179,16 +1179,41 @@ TYPED_TEST(SmallVectorReferenceInvalidationTest, AppendRange) {
 }
 
 TYPED_TEST(SmallVectorReferenceInvalidationTest, Assign) {
+  // 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
-  // Regardless of capacity, assign should never reference an internal element.
-  EXPECT_DEATH(V.assign(1, V.back()), this->AssertionMessage);
-  EXPECT_DEATH(V.assign(this->NumBuiltinElts(V), V.back()),
-               this->AssertionMessage);
-  EXPECT_DEATH(V.assign(this->NumBuiltinElts(V) + 1, V.back()),
-               this->AssertionMessage);
-#endif
+  int N = this->NumBuiltinElts(V);
+  ASSERT_EQ(unsigned(N), V.size());
+  ASSERT_EQ(unsigned(N), V.capacity());
+
+  // Check assign that shrinks in small mode.
+  V.assign(1, V.back());
+  EXPECT_EQ(1u, V.size());
+  EXPECT_EQ(N, V[0]);
+
+  // Check assign that grows within small mode.
+  ASSERT_LT(V.size(), V.capacity());
+  V.assign(V.capacity(), V.back());
+  for (int I = 0, E = V.size(); I != E; ++I) {
+    EXPECT_EQ(N, V[I]);
+
+    // Reset to [1, 2, ...].
+    V[I] = I + 1;
+  }
+
+  // Check assign that grows to large mode.
+  ASSERT_EQ(2, V[1]);
+  V.assign(V.capacity() + 1, V[1]);
+  for (int I = 0, E = V.size(); I != E; ++I) {
+    EXPECT_EQ(2, V[I]);
+
+    // Reset to [1, 2, ...].
+    V[I] = I + 1;
+  }
+
+  // Check assign that shrinks in large mode.
+  V.assign(1, V[1]);
+  EXPECT_EQ(2, V[0]);
 }
 
 TYPED_TEST(SmallVectorReferenceInvalidationTest, AssignRange) {
@@ -1289,11 +1314,25 @@ TYPED_TEST(SmallVectorReferenceInvalidationTest, InsertRange) {
 }
 
 TYPED_TEST(SmallVectorReferenceInvalidationTest, EmplaceBack) {
+  // 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.emplace_back(V.back()), this->AssertionMessage);
-#endif
+  int N = this->NumBuiltinElts(V);
+
+  // Push back a reference to last element when growing from small storage.
+  V.emplace_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.emplace_back(V.back());
+  EXPECT_EQ(int(V.size()) - 1, V.back());
 }
 
 template <class VectorT>
@@ -1315,25 +1354,42 @@ protected:
     SmallVectorTestBase::SetUp();
 
     // Fill up the small size so that insertions move the elements.
-    V.push_back(std::make_pair(0, 0));
+    for (int I = 0, E = NumBuiltinElts(V); I != E; ++I)
+      V.emplace_back(I + 1, I + 1);
   }
 };
 
 // Test pairs of the same types from SmallVectorReferenceInvalidationTestTypes.
 using SmallVectorInternalReferenceInvalidationTestTypes =
-    ::testing::Types<SmallVector<std::pair<int, int>, 1>,
-                     SmallVector<std::pair<Constructable, Constructable>, 1>>;
+    ::testing::Types<SmallVector<std::pair<int, int>, 3>,
+                     SmallVector<std::pair<Constructable, Constructable>, 3>>;
 
 TYPED_TEST_CASE(SmallVectorInternalReferenceInvalidationTest,
                 SmallVectorInternalReferenceInvalidationTestTypes);
 
 TYPED_TEST(SmallVectorInternalReferenceInvalidationTest, EmplaceBack) {
+  // 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.emplace_back(V.back().first, 0), this->AssertionMessage);
-  EXPECT_DEATH(V.emplace_back(0, V.back().second), this->AssertionMessage);
-#endif
+  int N = this->NumBuiltinElts(V);
+
+  // Push back a reference to last element when growing from small storage.
+  V.emplace_back(V.back().first, V.back().second);
+  EXPECT_EQ(N, V.back().first);
+  EXPECT_EQ(N, V.back().second);
+
+  // Check that the old value is still there (not moved away).
+  EXPECT_EQ(N, V[V.size() - 2].first);
+  EXPECT_EQ(N, V[V.size() - 2].second);
+
+  // Fill storage again.
+  V.back().first = V.back().second = V.size();
+  while (V.size() < V.capacity())
+    V.emplace_back(V.size() + 1, V.size() + 1);
+
+  // Push back a reference to last element when growing from large storage.
+  V.emplace_back(V.back().first, V.back().second);
+  EXPECT_EQ(int(V.size()) - 1, V.back().first);
+  EXPECT_EQ(int(V.size()) - 1, V.back().second);
 }
 
 } // end namespace