diff --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h --- a/llvm/include/llvm/ADT/SmallVector.h +++ b/llvm/include/llvm/ADT/SmallVector.h @@ -66,13 +66,27 @@ /// 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 @@ : public SmallVectorBase> { using Base = SmallVectorBase>; +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 @@ } // 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 @@ /// 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( - SmallVectorBase>::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 @@ takeAllocationForGrow(NewElts, NewCapacity); } +template +T *SmallVectorTemplateBase::mallocForGrow( + size_t MinSize, size_t &NewCapacity) { + return static_cast( + SmallVectorBase>::mallocForGrow( + this->getFirstEl(), MinSize, sizeof(T), NewCapacity)); +} + // Define this out-of-line to dissuade the C++ compiler from inlining it. template void SmallVectorTemplateBase::moveElementsForGrow( diff --git a/llvm/lib/Support/SmallVector.cpp b/llvm/lib/Support/SmallVector.cpp --- a/llvm/lib/Support/SmallVector.cpp +++ b/llvm/lib/Support/SmallVector.cpp @@ -108,12 +108,29 @@ return std::clamp(NewCapacity, MinSize, MaxSize); } +template +void *SmallVectorBase::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 -void *SmallVectorBase::mallocForGrow(size_t MinSize, size_t TSize, +void *SmallVectorBase::mallocForGrow(void *FirstEl, size_t MinSize, + size_t TSize, size_t &NewCapacity) { NewCapacity = getNewCapacity(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 @@ size_t NewCapacity = getNewCapacity(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;