ADT: Add assertions to SmallVector::insert, etc., for reference invalidation

[dexonsmith]

Summary

2c196bbc6bd897b3dcc1d87a3baac28e1e88df41 asserted that
SmallVector::push_back doesn't invalidate the parameter when it needs
to grow. Do the same for resize, append, assign, insert, and
emplace_back.

[mehdi_amini]

llvm/include/llvm/ADT/SmallVector.h

735

Do you need the SFINAE and the two implementation?
Wouldn't it always be valid to call this->assertSafeToAdd(&Arg);?

[dexonsmith]

Updates:

• Per @mehdi_amini's suggestion, check all arguments of emplace_back.
• Check ranges in the iterator range versions of assign, append, and insert.

[dexonsmith]

llvm/include/llvm/ADT/SmallVector.h

735

In fact we can also check for internal references in the multi-argument versions of emplace_back. The new patch does this... thanks!

[dexonsmith]

Running check-all for llvm/clang/clang-tools-extra hit three more assertions, now fixed (since splitting the patch out from https://reviews.llvm.org/D91467 I'd just run ADTTests until now):

• Two MCInst::insert calls.
• One emplace_back call with an internal reference.

I'll commit in a minute.

[dexonsmith]

I committed a follow-up for Windows at 90966daac3dd633b9621426e6621a706807d7c0d.

llvm/include/llvm/ADT/SmallVector.h

protected:
   bool isSmall() const { return this->BeginX == getFirstEl(); }
 
   /// Put this vector in a state of being small.
   void resetToSmall() {
     this->BeginX = getFirstEl();
     this->Size = this->Capacity = 0; // FIXME: Setting Capacity to 0 is suspect.
   }
 
-  void assertSafeToPush(const void *Elt) {
-    assert(
-        (Elt < begin() || Elt >= end() || this->size() < this->capacity()) &&
-        "Attempting to push_back to the vector an element of the vector without"
-        " enough space reserved");
+  /// Check whether Elt will be invalidated by resizing the vector to NewSize.
+  void assertSafeToReferenceAfterResize(const void *Elt, size_t NewSize) {
+    assert((Elt >= this->end() ||
+            (NewSize <= this->size()
+                 ? Elt < this->begin() + NewSize
+                 : (Elt < this->begin() || NewSize <= this->capacity()))) &&
+           "Attempting to reference an element of the vector in an operation "
+           "that invalidates it");
+  }
+
+  /// Check whether Elt will be invalidated by increasing the size of the
+  /// vector by N.
+  void assertSafeToAdd(const void *Elt, size_t N = 1) {
+    return this->assertSafeToReferenceAfterResize(Elt, this->size() + N);
   }
 
 public:
   using size_type = size_t;
   using difference_type = ptrdiff_t;
   using value_type = T;
   using iterator = T *;
   using const_iterator = const 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);
 
 public:
   void push_back(const T &Elt) {
-    this->assertSafeToPush(&Elt);
+    this->assertSafeToAdd(&Elt);
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
     ::new ((void*) this->end()) T(Elt);
     this->set_size(this->size() + 1);
   }
 
   void push_back(T &&Elt) {
-    this->assertSafeToPush(&Elt);
+    this->assertSafeToAdd(&Elt);
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
     ::new ((void*) this->end()) T(::std::move(Elt));
     this->set_size(this->size() + 1);
   }
 
   void pop_back() {
     this->set_size(this->size() - 1);
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)); }
 
 public:
   void push_back(const T &Elt) {
-    this->assertSafeToPush(&Elt);
+    this->assertSafeToAdd(&Elt);
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
     memcpy(reinterpret_cast<void *>(this->end()), &Elt, sizeof(T));
     this->set_size(this->size() + 1);
   }
 
   void pop_back() { this->set_size(this->size() - 1); }
 };
if (N < this->size()) {
         this->grow(N);
       for (auto I = this->end(), E = this->begin() + N; I != E; ++I)
         new (&*I) T();
       this->set_size(N);
     }
   }
 
   void resize(size_type N, const T &NV) {
+    this->assertSafeToReferenceAfterResize(&NV, N);
     if (N < this->size()) {
       this->destroy_range(this->begin()+N, this->end());
       this->set_size(N);
     } else if (N > this->size()) {
       if (this->capacity() < N)
         this->grow(N);
       std::uninitialized_fill(this->end(), this->begin()+N, NV);
       this->set_size(N);
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);
     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 &
   // re-initializing them - for all assign(...) variants.
 
   void assign(size_type NumElts, const T &Elt) {
+    this->assertSafeToReferenceAfterResize(&Elt, 0);
     clear();
     if (this->capacity() < NumElts)
       this->grow(NumElts);
     this->set_size(NumElts);
     std::uninitialized_fill(this->begin(), this->end(), Elt);
   }
 
   template <typename in_iter,
template <class ArgType> iterator insert_one_impl(iterator I, ArgType &&Elt) {
     if (I == this->end()) {  // Important special case for empty vector.
       this->push_back(::std::forward<ArgType>(Elt));
       return this->end()-1;
     }
 
     assert(I >= this->begin() && "Insertion iterator is out of bounds.");
     assert(I <= this->end() && "Inserting past the end of the vector.");
 
+    // 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;
     }
 
     ::new ((void*) this->end()) T(::std::move(this->back()));
     // Push everything else over.
iterator insert(iterator I, size_type NumToInsert, const T &Elt) {
     if (I == this->end()) {  // Important special case for empty vector.
       append(NumToInsert, Elt);
       return this->begin()+InsertElt;
     }
 
     assert(I >= this->begin() && "Insertion iterator is out of bounds.");
     assert(I <= this->end() && "Inserting past the end of the vector.");
 
+    // Check that adding NumToInsert elements won't invalidate Elt.
+    this->assertSafeToAdd(&Elt, NumToInsert);
+
     // Ensure there is enough space.
     reserve(this->size() + 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
iterator insert(iterator I, ItTy From, ItTy To) {
     this->uninitialized_copy(From, To, OldEnd);
     return I;
   }
 
   void insert(iterator I, std::initializer_list<T> IL) {
     insert(I, IL.begin(), IL.end());
   }
 
-  template <typename... ArgTypes> reference emplace_back(ArgTypes &&... Args) {
+private:
+  template <class... ArgTypes>
+  reference emplace_back_impl(ArgTypes &&... Args) {
     if (LLVM_UNLIKELY(this->size() >= this->capacity()))
       this->grow();
     ::new ((void *)this->end()) T(std::forward<ArgTypes>(Args)...);
     this->set_size(this->size() + 1);
     return this->back();
   }
 
+public:
+  reference emplace_back() { return emplace_back_impl(); }
+
+  template <
+      class ArgType,
+      std::enable_if_t<
+          !std::is_same<std::remove_const_t<std::remove_reference_t<ArgType>>,
+                        T>::value,
+          bool> = true>
+  reference emplace_back(ArgType &&Arg) {
+    return emplace_back_impl(std::forward<ArgType>(Arg));
+  }
+
+  template <
+      class ArgType,
+      std::enable_if_t<
+          std::is_same<std::remove_const_t<std::remove_reference_t<ArgType>>,
+                       T>::value,
+          bool> = true>
+  reference emplace_back(ArgType &&Arg) {
+    // ArgType is a T and may be a reference into this vector. Check whether
+    // it's safe before forwarding to emplace_back_impl.
+    this->assertSafeToAdd(&Arg);
+    return emplace_back_impl(std::forward<ArgType>(Arg));
+  }
+
+  template <class ArgType1, class ArgType2, class... ArgTypes>
+  reference emplace_back(ArgType1 &&Arg1, ArgType2 &&Arg2,
+                         ArgTypes &&... Args) {
+    return emplace_back_impl(std::forward<ArgType1>(Arg1),
+                             std::forward<ArgType2>(Arg2),
+                             std::forward<ArgTypes>(Args)...);
+  }
+
   SmallVectorImpl &operator=(const SmallVectorImpl &RHS);
 
   SmallVectorImpl &operator=(SmallVectorImpl &&RHS);
 
   bool operator==(const SmallVectorImpl &RHS) const {
     if (this->size() != RHS.size()) return false;
     return std::equal(this->begin(), this->end(), RHS.begin());
   }

llvm/unittests/ADT/SmallVectorTest.cpp

TYPED_TEST(SmallVectorTest, PushPopTest) {
   this->assertValuesInOrder(this->theVector, 1u, 1);
   EXPECT_FALSE(this->theVector.begin() == this->theVector.end());
   EXPECT_FALSE(this->theVector.empty());
 
   // Push another element
   this->theVector.push_back(Constructable(2));
   this->assertValuesInOrder(this->theVector, 2u, 1, 2);
 
-  // Insert at beginning
+  // Insert at beginning. Reserve space to avoid reference invalidation from
+  // this->theVector[1].
+  this->theVector.reserve(this->theVector.size() + 1);
   this->theVector.insert(this->theVector.begin(), this->theVector[1]);
   this->assertValuesInOrder(this->theVector, 3u, 2, 1, 2);
 
   // Pop one element
   this->theVector.pop_back();
   this->assertValuesInOrder(this->theVector, 2u, 2, 1);
 
   // Pop remaining elements
TEST(SmallVectorTest, InitializerList) {
   V2.assign({4});
   EXPECT_TRUE(makeArrayRef(V2).equals({4}));
   V2.append({3, 2});
   EXPECT_TRUE(makeArrayRef(V2).equals({4, 3, 2}));
   V2.insert(V2.begin() + 1, 5);
   EXPECT_TRUE(makeArrayRef(V2).equals({4, 5, 3, 2}));
 }
 
+template <class VectorT>
+class SmallVectorReferenceInvalidationTest : public SmallVectorTestBase {
+protected:
+  const char *AssertionMessage =
+      "Attempting to reference an element of the vector in an operation \" "
+      "\"that invalidates it";
+
+  VectorT V;
+
+  template <typename T, unsigned N>
+  static unsigned NumBuiltinElts(const SmallVector<T, N> &) {
+    return N;
+  }
+
+  void SetUp() override {
+    SmallVectorTestBase::SetUp();
+
+    // Fill up the small size so that insertions move the elements.
+    V.push_back(0);
+  }
+};
+
+// Test one type that's trivially copyable (int) and one that isn't
+// (Constructable) since reference invalidation may be fixed differently for
+// each.
+using SmallVectorReferenceInvalidationTestTypes =
+    ::testing::Types<SmallVector<int, 1>, SmallVector<Constructable, 1>>;
+
+TYPED_TEST_CASE(SmallVectorReferenceInvalidationTest,
+                SmallVectorReferenceInvalidationTestTypes);
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, PushBack) {
+  auto &V = this->V;
+  (void)V;
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(V.push_back(V.back()), this->AssertionMessage);
+#endif
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, PushBackMoved) {
+  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
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, Resize) {
+  auto &V = this->V;
+  (void)V;
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(V.resize(2, V.back()), this->AssertionMessage);
+#endif
+}
+
+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
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, Assign) {
+  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(2, V.back()), this->AssertionMessage);
+  EXPECT_DEATH(V.assign(3, V.back()), this->AssertionMessage);
+#endif
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, Insert) {
+  auto &V = this->V;
+  (void)V;
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(V.insert(V.begin(), V.back()), this->AssertionMessage);
+#endif
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, InsertMoved) {
+  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
+}
+
+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
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, EmplaceBack) {
+  auto &V = this->V;
+  (void)V;
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(V.emplace_back(V.back()), this->AssertionMessage);
+#endif
+}
+
+TYPED_TEST(SmallVectorReferenceInvalidationTest, EmplaceBackMoved) {
+  auto &V = this->V;
+  (void)V;
+#if !defined(NDEBUG) && GTEST_HAS_DEATH_TEST
+  EXPECT_DEATH(V.emplace_back(std::move(V.back())), this->AssertionMessage);
+#endif
+}
+
 } // end namespace

llvm/utils/TableGen/CodeGenSchedule.cpp

pushVariant(const TransVariant &VInfo, bool IsRead) {
 
   const CodeGenSchedRW &SchedRW = SchedModels.getSchedRW(VInfo.RWIdx, IsRead);
 
   SmallVectorImpl<SmallVector<unsigned,4>> &RWSequences = IsRead
     ? Trans.ReadSequences : Trans.WriteSequences;
   if (SchedRW.IsVariadic) {
     unsigned OperIdx = RWSequences.size()-1;
     // Make N-1 copies of this transition's last sequence.
+    RWSequences.reserve(RWSequences.size() + SelectedRWs.size() - 1);
     RWSequences.insert(RWSequences.end(), SelectedRWs.size() - 1,
                        RWSequences[OperIdx]);
     // Push each of the N elements of the SelectedRWs onto a copy of the last
     // sequence (split the current operand into N operands).
     // Note that write sequences should be expanded within this loop--the entire
     // sequence belongs to a single operand.
     for (IdxIter RWI = SelectedRWs.begin(), RWE = SelectedRWs.end();
          RWI != RWE; ++RWI, ++OperIdx) {
if (IsRead && IsForAnyCPU && HasAliases && !WasPushed) {
       // a) We have read variant for CPU A
       // b) We have write variant for CPU B
       // b) We don't have write variant for CPU A
       // d) We must expand all read/write variants (IsForAnyCPU is true)
       // e) We couldn't expand SchedRW because TransVec doesn't have
       //    any transition with compatible CPU ID.
       // In such case we create new empty transition with zero (AnyCPU)
       // index.
+      TransVec.reserve(TransVec.size() + 1);
       TransVec.emplace_back(TransVec[StartIdx].PredTerm);
       TransVec.back().ReadSequences.emplace_back();
       CollectAndAddVariants(TransVec.size() - 1, SchedRW);
     }
   }
 }
 
 // For each variant of a Read/Write in Trans, substitute the sequence of