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

[Alina Sbirlea via llvm-commits]

Author: Alina Sbirlea
Date: 2022-10-06T15:37:43-07:00
New Revision: 0fbeca0ee6f0cdc37b89e5b06fbb3ca6caf26f66

URL: https://github.com/llvm/llvm-project/commit/0fbeca0ee6f0cdc37b89e5b06fbb3ca6caf26f66
DIFF: https://github.com/llvm/llvm-project/commit/0fbeca0ee6f0cdc37b89e5b06fbb3ca6caf26f66.diff

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

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

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

Added: 
    

Modified: 
    llvm/include/llvm/ADT/SmallVector.h
    llvm/lib/Support/SmallVector.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h
index 513fce93549fb..bbc5bc8ceeea8 100644
--- a/llvm/include/llvm/ADT/SmallVector.h
+++ b/llvm/include/llvm/ADT/SmallVector.h
@@ -66,13 +66,27 @@ template <class Size_T> class SmallVectorBase {
   /// 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; }
@@ -110,6 +124,7 @@ 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.
@@ -120,7 +135,6 @@ class SmallVectorTemplateCommon
   }
   // 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) {
@@ -352,11 +366,7 @@ class SmallVectorTemplateBase : public SmallVectorTemplateCommon<T> {
 
   /// 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().
@@ -430,6 +440,14 @@ void SmallVectorTemplateBase<T, TriviallyCopyable>::grow(size_t MinSize) {
   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(

diff  --git a/llvm/lib/Support/SmallVector.cpp b/llvm/lib/Support/SmallVector.cpp
index 062e96544734a..f7e7e80332cc3 100644
--- a/llvm/lib/Support/SmallVector.cpp
+++ b/llvm/lib/Support/SmallVector.cpp
@@ -108,12 +108,29 @@ static size_t getNewCapacity(size_t MinSize, size_t TSize, size_t OldCapacity) {
   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.
@@ -123,13 +140,17 @@ void SmallVectorBase<Size_T>::grow_pod(void *FirstEl, size_t MinSize,
   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;