[SmallVector] Reallocate if assigned memory is right after the current vector, created with capacity 0

[asbirlea]

Summary

Potential solution for
https://github.com/llvm/llvm-project/issues/57324.

[dblaikie]

If possible, it'd be good to have some performance data on this since it adds an extra condition on every vector resize.

[asbirlea]

Compile-time impact looks neutral:
https://llvm-compile-time-tracker.com/compare.php?from=8d5b694da172c97f4fc10d5586830fea1b486037&to=00e94b5581ec2361ca2261770addd700fcf49574&stat=instructions

[fhahn]

Can this be limited to instantiations where the original original size is 0, as IIUC this only applies to llvm::SmallVector<T, 0>?

[asbirlea]

In D132512#3758117, @fhahn wrote:

Can this be limited to instantiations where the original original size is 0, as IIUC this only applies to llvm::SmallVector<T, 0>?

I don't know how to properly check for that.

Say the original instantiation is llvm::SmallVector<T, 0>, so initially capacity is set to 0, but after the vector is used, the size and capacity have increased, and the heap allocated memory is in a completely different region (safe). Now say that we're at a point where we want to further increase capacity, and it just so happens that the memory right after (at FirstEl) is available and, this time, *that's* what safe_realloc returns. This is the case at line:153.
The case at line :145, we're already in a case where the vector has no heap allocated since BeginX == FirstEl (if we entered this condition and there's heap allocated, we failed to fix the problem this patch is trying to fix...). So the only scenario in which NewElts can be equal to FirstEl is when the vector was created with capacity 0. The check would be redundant here.
The third case, when mallocForGrow is called (line :132), it's also possible the capacity was already increased and this is called to increase it further. So checking for capacity 0 here wouldn't stop allocating right at FirstEl.

Please let me know if I missed something in the above reasoning.

[asbirlea]

ping.

[dblaikie]

Seems good to me.

llvm/include/llvm/ADT/SmallVector.h

protected:
 
   SmallVectorBase() = delete;
   SmallVectorBase(void *FirstEl, size_t TotalCapacity)
       : BeginX(FirstEl), Capacity(TotalCapacity) {}
 
   /// This is a helper for \a grow() that's out of line to reduce code
   /// duplication.  This function will report a fatal error if it can't grow at
   /// least to \p MinSize.
-  void *mallocForGrow(size_t MinSize, size_t TSize, size_t &NewCapacity);
+  void *mallocForGrow(void *FirstEl, size_t MinSize, size_t TSize,
+                      size_t &NewCapacity);
 
   /// This is an implementation of the grow() method which only works
   /// on POD-like data types and is out of line to reduce code duplication.
   /// This function will report a fatal error if it cannot increase capacity.
   void grow_pod(void *FirstEl, size_t MinSize, size_t TSize);
 
+  /// If vector was first created with capacity 0, getFirstEl() points to the
+  /// memory right after, an area unallocated. If a subsequent allocation,
+  /// that grows the vector, happens to return the same pointer as getFirstEl(),
+  /// get a new allocation, otherwise isSmall() will falsely return that no
+  /// allocation was done (true) and the memory will not be freed in the
+  /// destructor. If a VSize is given (vector size), also copy that many
+  /// elements to the new allocation - used if realloca fails to increase
+  /// space, and happens to allocate precisely at BeginX.
+  /// This is unlikely to be called often, but resolves a memory leak when the
+  /// situation does occur.
+  void *replaceAllocation(void *NewElts, size_t TSize, size_t NewCapacity,
+                          size_t VSize = 0);
+
 public:
   size_t size() const { return Size; }
   size_t capacity() const { return Capacity; }
 
   [[nodiscard]] bool empty() const { return !Size; }
 
 protected:
   /// Set the array size to \p N, which the current array must have enough

 /// This is the part of SmallVectorTemplateBase which does not depend on whether
 /// the type T is a POD. The extra dummy template argument is used by ArrayRef
 /// to avoid unnecessarily requiring T to be complete.
 template <typename T, typename = void>
 class SmallVectorTemplateCommon
     : public SmallVectorBase<SmallVectorSizeType<T>> {
   using Base = SmallVectorBase<SmallVectorSizeType<T>>;
 
+protected:
   /// Find the address of the first element.  For this pointer math to be valid
   /// with small-size of 0 for T with lots of alignment, it's important that
   /// SmallVectorStorage is properly-aligned even for small-size of 0.
   void *getFirstEl() const {
     return const_cast<void *>(reinterpret_cast<const void *>(
         reinterpret_cast<const char *>(this) +
         offsetof(SmallVectorAlignmentAndSize<T>, FirstEl)));
   }
   // Space after 'FirstEl' is clobbered, do not add any instance vars after it.
 
-protected:
   SmallVectorTemplateCommon(size_t Size) : Base(getFirstEl(), Size) {}
 
   void grow_pod(size_t MinSize, size_t TSize) {
     Base::grow_pod(getFirstEl(), MinSize, TSize);
   }
 
   /// Return true if this is a smallvector which has not had dynamic
   /// memory allocated for it.
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);
 
   /// Create a new allocation big enough for \p MinSize and pass back its size
   /// in \p NewCapacity. This is the first section of \a grow().
-  T *mallocForGrow(size_t MinSize, size_t &NewCapacity) {
-    return static_cast<T *>(
-        SmallVectorBase<SmallVectorSizeType<T>>::mallocForGrow(
-            MinSize, sizeof(T), NewCapacity));
-  }
+  T *mallocForGrow(size_t MinSize, size_t &NewCapacity);
 
   /// Move existing elements over to the new allocation \p NewElts, the middle
   /// section of \a grow().
   void moveElementsForGrow(T *NewElts);
 
   /// Transfer ownership of the allocation, finishing up \a grow().
   void takeAllocationForGrow(T *NewElts, size_t NewCapacity);
 

 template <typename T, bool TriviallyCopyable>
 void SmallVectorTemplateBase<T, TriviallyCopyable>::grow(size_t MinSize) {
   size_t NewCapacity;
   T *NewElts = mallocForGrow(MinSize, NewCapacity);
   moveElementsForGrow(NewElts);
   takeAllocationForGrow(NewElts, NewCapacity);
 }
 
+template <typename T, bool TriviallyCopyable>
+T *SmallVectorTemplateBase<T, TriviallyCopyable>::mallocForGrow(
+    size_t MinSize, size_t &NewCapacity) {
+  return static_cast<T *>(
+      SmallVectorBase<SmallVectorSizeType<T>>::mallocForGrow(
+          this->getFirstEl(), MinSize, sizeof(T), NewCapacity));
+}
+
 // Define this out-of-line to dissuade the C++ compiler from inlining it.
 template <typename T, bool TriviallyCopyable>
 void SmallVectorTemplateBase<T, TriviallyCopyable>::moveElementsForGrow(
     T *NewElts) {
   // Move the elements over.
   this->uninitialized_move(this->begin(), this->end(), NewElts);
 
   // Destroy the original elements.

llvm/lib/Support/SmallVector.cpp

if (OldCapacity == MaxSize)
     report_at_maximum_capacity(MaxSize);
 
   // In theory 2*capacity can overflow if the capacity is 64 bit, but the
   // original capacity would never be large enough for this to be a problem.
   size_t NewCapacity = 2 * OldCapacity + 1; // Always grow.
   return std::clamp(NewCapacity, MinSize, MaxSize);
 }
 
+template <class Size_T>
+void *SmallVectorBase<Size_T>::replaceAllocation(void *NewElts, size_t TSize,
+                                                 size_t NewCapacity,
+                                                 size_t VSize) {
+  void *NewEltsReplace = llvm::safe_malloc(NewCapacity * TSize);
+  if (VSize)
+    memcpy(NewEltsReplace, NewElts, VSize * TSize);
+  free(NewElts);
+  return NewEltsReplace;
+}
+
 // Note: Moving this function into the header may cause performance regression.
 template <class Size_T>
-void *SmallVectorBase<Size_T>::mallocForGrow(size_t MinSize, size_t TSize,
+void *SmallVectorBase<Size_T>::mallocForGrow(void *FirstEl, size_t MinSize,
+                                             size_t TSize,
                                              size_t &NewCapacity) {
   NewCapacity = getNewCapacity<Size_T>(MinSize, TSize, this->capacity());
-  return llvm::safe_malloc(NewCapacity * TSize);
+  // Even if capacity is not 0 now, if the vector was originally created with
+  // capacity 0, it's possible for the malloc to return FirstEl.
+  void *NewElts = llvm::safe_malloc(NewCapacity * TSize);
+  if (NewElts == FirstEl)
+    NewElts = replaceAllocation(NewElts, TSize, NewCapacity);
+  return NewElts;
 }
 
 // Note: Moving this function into the header may cause performance regression.
 template <class Size_T>
 void SmallVectorBase<Size_T>::grow_pod(void *FirstEl, size_t MinSize,
                                        size_t TSize) {
   size_t NewCapacity = getNewCapacity<Size_T>(MinSize, TSize, this->capacity());
   void *NewElts;
   if (BeginX == FirstEl) {
-    NewElts = safe_malloc(NewCapacity * TSize);
+    NewElts = llvm::safe_malloc(NewCapacity * TSize);
+    if (NewElts == FirstEl)
+      NewElts = replaceAllocation(NewElts, TSize, NewCapacity);
 
     // Copy the elements over.  No need to run dtors on PODs.
     memcpy(NewElts, this->BeginX, size() * TSize);
   } else {
     // If this wasn't grown from the inline copy, grow the allocated space.
-    NewElts = safe_realloc(this->BeginX, NewCapacity * TSize);
+    NewElts = llvm::safe_realloc(this->BeginX, NewCapacity * TSize);
+    if (NewElts == FirstEl)
+      NewElts = replaceAllocation(NewElts, TSize, NewCapacity, size());
   }
 
   this->BeginX = NewElts;
   this->Capacity = NewCapacity;
 }
 
 template class llvm::SmallVectorBase<uint32_t>;