[Support] [IR] [ADT] - Check nullptr after allocation with malloc/realloc or calloc

[kkretzsch]

Summary

As a follow up of the bad alloc handler patch, this patch introduces nullptr checks on pointers returned from the
malloc/realloc/calloc functions. In addition some memory size assignments are moved behind the allocation
of the corresponding memory to fulfill exception safe memory management (RAII).

patch by Klaus Kretzschmar

[MatzeB]

LGTM

[kkretzsch]

Thanks for reviewing my change. Since I do not have commit access to trunk, can somebody push my change to trunk?

Thanks!
Klaus

[MatzeB]

Finally got around pushing this. You may consider asking for commit access (https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access) if you plan further changes.

[kkretzsch]

In D35414#815546, @MatzeB wrote:

Finally got around pushing this. You may consider asking for commit access (https://llvm.org/docs/DeveloperPolicy.html#obtaining-commit-access) if you plan further changes.

Thanks Matthias! Yes, I'll ask for the commit access.

llvm/trunk/include/llvm/ADT/SmallVector.h

 #ifndef LLVM_ADT_SMALLVECTOR_H
 #define LLVM_ADT_SMALLVECTOR_H
 
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/AlignOf.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/type_traits.h"
+#include "llvm/Support/ErrorHandling.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
 #include <cstdlib>
 #include <cstring>
 #include <initializer_list>
 #include <iterator>
 #include <memory>

 void SmallVectorTemplateBase<T, isPodLike>::grow(size_t MinSize) {
   size_t CurCapacity = this->capacity();
   size_t CurSize = this->size();
   // Always grow, even from zero.
   size_t NewCapacity = size_t(NextPowerOf2(CurCapacity+2));
   if (NewCapacity < MinSize)
     NewCapacity = MinSize;
   T *NewElts = static_cast<T*>(malloc(NewCapacity*sizeof(T)));
+  if (NewElts == nullptr)
+    report_bad_alloc_error("Allocation of SmallVector element failed.");
 
   // Move the elements over.
   this->uninitialized_move(this->begin(), this->end(), NewElts);
 
   // Destroy the original elements.
   destroy_range(this->begin(), this->end());
 
   // If this wasn't grown from the inline copy, deallocate the old space.

llvm/trunk/include/llvm/ADT/StringMap.h

 #ifndef LLVM_ADT_STRINGMAP_H
 #define LLVM_ADT_STRINGMAP_H
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
+#include "llvm/Support/ErrorHandling.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
 #include <initializer_list>
 #include <iterator>
 #include <utility>
static StringMapEntry *Create(StringRef Key, AllocatorTy &Allocator,
     // terminator.
     unsigned AllocSize = static_cast<unsigned>(sizeof(StringMapEntry))+
       KeyLength+1;
     unsigned Alignment = alignof(StringMapEntry);
 
     StringMapEntry *NewItem =
       static_cast<StringMapEntry*>(Allocator.Allocate(AllocSize,Alignment));
 
+    if (NewItem == nullptr)
+      report_bad_alloc_error("Allocation of StringMap entry failed.");
+
     // Construct the value.
     new (NewItem) StringMapEntry(KeyLength, std::forward<InitTy>(InitVals)...);
 
     // Copy the string information.
     char *StrBuffer = const_cast<char*>(NewItem->getKeyData());
     if (KeyLength > 0)
       memcpy(StrBuffer, Key.data(), KeyLength);
     StrBuffer[KeyLength] = 0;  // Null terminate for convenience of clients.

llvm/trunk/lib/IR/DataLayout.cpp

const StructLayout *DataLayout::getStructLayout(StructType *Ty) const {
   StructLayout *&SL = (*STM)[Ty];
   if (SL) return SL;
 
   // Otherwise, create the struct layout.  Because it is variable length, we
   // malloc it, then use placement new.
   int NumElts = Ty->getNumElements();
   StructLayout *L =
     (StructLayout *)malloc(sizeof(StructLayout)+(NumElts-1) * sizeof(uint64_t));
+  if (L == nullptr)
+    report_bad_alloc_error("Allocation of StructLayout elements failed.");
 
   // Set SL before calling StructLayout's ctor.  The ctor could cause other
   // entries to be added to TheMap, invalidating our reference.
   SL = L;
 
   new (L) StructLayout(Ty, *this);
 
   return L;

llvm/trunk/lib/Support/FoldingSet.cpp

static void **GetBucketFor(unsigned Hash, void **Buckets, unsigned NumBuckets) {
   // NumBuckets is always a power of 2.
   unsigned BucketNum = Hash & (NumBuckets-1);
   return Buckets + BucketNum;
 }
 
 /// AllocateBuckets - Allocated initialized bucket memory.
 static void **AllocateBuckets(unsigned NumBuckets) {
   void **Buckets = static_cast<void**>(calloc(NumBuckets+1, sizeof(void*)));
+
+  if (Buckets == nullptr)
+    report_bad_alloc_error("Allocation of Buckets failed.");
+  
   // Set the very last bucket to be a non-null "pointer".
   Buckets[NumBuckets] = reinterpret_cast<void*>(-1);
   return Buckets;
 }
 
 //===----------------------------------------------------------------------===//
 // FoldingSetBase Implementation
 
void FoldingSetBase::clear() {
   NumNodes = 0;
 }
 
 void FoldingSetBase::GrowBucketCount(unsigned NewBucketCount) {
   assert((NewBucketCount > NumBuckets) && "Can't shrink a folding set with GrowBucketCount");
   assert(isPowerOf2_32(NewBucketCount) && "Bad bucket count!");
   void **OldBuckets = Buckets;
   unsigned OldNumBuckets = NumBuckets;
-  NumBuckets = NewBucketCount;
   
   // Clear out new buckets.
-  Buckets = AllocateBuckets(NumBuckets);
+  Buckets = AllocateBuckets(NewBucketCount);
+  // Set NumBuckets only if allocation of new buckets was succesful
+  NumBuckets = NewBucketCount; 
   NumNodes = 0;
 
   // Walk the old buckets, rehashing nodes into their new place.
   FoldingSetNodeID TempID;
   for (unsigned i = 0; i != OldNumBuckets; ++i) {
     void *Probe = OldBuckets[i];
     if (!Probe) continue;
     while (Node *NodeInBucket = GetNextPtr(Probe)) {

llvm/trunk/lib/Support/SmallPtrSet.cpp

 // This file implements the SmallPtrSet class.  See SmallPtrSet.h for an
 // overview of the algorithm.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/ErrorHandling.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdlib>
 
 using namespace llvm;
 
 void SmallPtrSetImplBase::shrink_and_clear() {
   assert(!isSmall() && "Can't shrink a small set!");
   free(CurArray);
 
   // Reduce the number of buckets.
   unsigned Size = size();
   CurArraySize = Size > 16 ? 1 << (Log2_32_Ceil(Size) + 1) : 32;
   NumNonEmpty = NumTombstones = 0;
 
   // Install the new array.  Clear all the buckets to empty.
   CurArray = (const void**)malloc(sizeof(void*) * CurArraySize);
-  assert(CurArray && "Failed to allocate memory?");
+  if (CurArray == nullptr)
+    report_bad_alloc_error("Allocation of SmallPtrSet bucket array failed.");
+
   memset(CurArray, -1, CurArraySize*sizeof(void*));
 }
 
 std::pair<const void *const *, bool>
 SmallPtrSetImplBase::insert_imp_big(const void *Ptr) {
   if (LLVM_UNLIKELY(size() * 4 >= CurArraySize * 3)) {
     // If more than 3/4 of the array is full, grow.
     Grow(CurArraySize < 64 ? 128 : CurArraySize * 2);

 /// Grow - Allocate a larger backing store for the buckets and move it over.
 ///
 void SmallPtrSetImplBase::Grow(unsigned NewSize) {
   const void **OldBuckets = CurArray;
   const void **OldEnd = EndPointer();
   bool WasSmall = isSmall();
 
   // Install the new array.  Clear all the buckets to empty.
-  CurArray = (const void**)malloc(sizeof(void*) * NewSize);
-  assert(CurArray && "Failed to allocate memory?");
+  const void **NewBuckets = (const void**) malloc(sizeof(void*) * NewSize);
+  if (NewBuckets == nullptr)
+    report_bad_alloc_error("Allocation of SmallPtrSet bucket array failed.");
+
+  // Reset member only if memory was allocated successfully
+  CurArray = NewBuckets;
   CurArraySize = NewSize;
   memset(CurArray, -1, NewSize*sizeof(void*));
 
   // Copy over all valid entries.
   for (const void **BucketPtr = OldBuckets; BucketPtr != OldEnd; ++BucketPtr) {
     // Copy over the element if it is valid.
     const void *Elt = *BucketPtr;
     if (Elt != getTombstoneMarker() && Elt != getEmptyMarker())
SmallPtrSetImplBase::SmallPtrSetImplBase(const void **SmallStorage,
   SmallArray = SmallStorage;
 
   // If we're becoming small, prepare to insert into our stack space
   if (that.isSmall()) {
     CurArray = SmallArray;
   // Otherwise, allocate new heap space (unless we were the same size)
   } else {
     CurArray = (const void**)malloc(sizeof(void*) * that.CurArraySize);
-    assert(CurArray && "Failed to allocate memory?");
+    if (CurArray == nullptr)
+      report_bad_alloc_error("Allocation of SmallPtrSet bucket array failed.");
   }
 
   // Copy over the that array.
   CopyHelper(that);
 }
 
 SmallPtrSetImplBase::SmallPtrSetImplBase(const void **SmallStorage,
                                          unsigned SmallSize,
if (isSmall())
       CurArray = (const void**)malloc(sizeof(void*) * RHS.CurArraySize);
     else {
       const void **T = (const void**)realloc(CurArray,
                                              sizeof(void*) * RHS.CurArraySize);
       if (!T)
         free(CurArray);
       CurArray = T;
     }
-    assert(CurArray && "Failed to allocate memory?");
+    if (CurArray == nullptr)
+      report_bad_alloc_error("Allocation of SmallPtrSet bucket array failed.");
   }
 
   CopyHelper(RHS);
 }
 
 void SmallPtrSetImplBase::CopyHelper(const SmallPtrSetImplBase &RHS) {
   // Copy over the new array size
   CurArraySize = RHS.CurArraySize;

llvm/trunk/lib/Support/SmallVector.cpp

void SmallVectorBase::grow_pod(void *FirstEl, size_t MinSizeInBytes,
   size_t CurSizeBytes = size_in_bytes();
   size_t NewCapacityInBytes = 2 * capacity_in_bytes() + TSize; // Always grow.
   if (NewCapacityInBytes < MinSizeInBytes)
     NewCapacityInBytes = MinSizeInBytes;
 
   void *NewElts;
   if (BeginX == FirstEl) {
     NewElts = malloc(NewCapacityInBytes);
+    if (NewElts == nullptr)
+      report_bad_alloc_error("Allocation of SmallVector element failed.");
 
     // Copy the elements over.  No need to run dtors on PODs.
     memcpy(NewElts, this->BeginX, CurSizeBytes);
   } else {
     // If this wasn't grown from the inline copy, grow the allocated space.
     NewElts = realloc(this->BeginX, NewCapacityInBytes);
+    if (NewElts == nullptr)
+      report_bad_alloc_error("Reallocation of SmallVector element failed.");
   }
-  assert(NewElts && "Out of memory");
 
   this->EndX = (char*)NewElts+CurSizeBytes;
   this->BeginX = NewElts;
   this->CapacityX = (char*)this->BeginX + NewCapacityInBytes;
 }

llvm/trunk/lib/Support/StringMap.cpp

StringMapImpl::StringMapImpl(unsigned InitSize, unsigned itemSize) {
   NumBuckets = 0;
   NumItems = 0;
   NumTombstones = 0;
 }
 
 void StringMapImpl::init(unsigned InitSize) {
   assert((InitSize & (InitSize-1)) == 0 &&
          "Init Size must be a power of 2 or zero!");
-  NumBuckets = InitSize ? InitSize : 16;
+
+  unsigned NewNumBuckets = InitSize ? InitSize : 16;
   NumItems = 0;
   NumTombstones = 0;
   
-  TheTable = (StringMapEntryBase **)calloc(NumBuckets+1,
+  TheTable = (StringMapEntryBase **)calloc(NewNumBuckets+1,
                                            sizeof(StringMapEntryBase **) +
                                            sizeof(unsigned));
 
+  if (TheTable == nullptr)
+    report_bad_alloc_error("Allocation of StringMap table failed.");
+
+  // Set the member only if TheTable was successfully allocated
+  NumBuckets = NewNumBuckets;
+
   // Allocate one extra bucket, set it to look filled so the iterators stop at
   // end.
   TheTable[NumBuckets] = (StringMapEntryBase*)2;
 }
 
 /// LookupBucketFor - Look up the bucket that the specified string should end
 /// up in.  If it already exists as a key in the map, the Item pointer for the
 /// specified bucket will be non-null.  Otherwise, it will be null.  In either
unsigned StringMapImpl::RehashTable(unsigned BucketNo) {
   }
 
   unsigned NewBucketNo = BucketNo;
   // Allocate one extra bucket which will always be non-empty.  This allows the
   // iterators to stop at end.
   StringMapEntryBase **NewTableArray =
     (StringMapEntryBase **)calloc(NewSize+1, sizeof(StringMapEntryBase *) +
                                              sizeof(unsigned));
+
+  if (NewTableArray == nullptr)
+    report_bad_alloc_error("Allocation of StringMap hash table failed.");
+
   unsigned *NewHashArray = (unsigned *)(NewTableArray + NewSize + 1);
   NewTableArray[NewSize] = (StringMapEntryBase*)2;
 
   // Rehash all the items into their new buckets.  Luckily :) we already have
   // the hash values available, so we don't have to rehash any strings.
   for (unsigned I = 0, E = NumBuckets; I != E; ++I) {
     StringMapEntryBase *Bucket = TheTable[I];
     if (Bucket && Bucket != getTombstoneVal()) {