[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.

include/llvm/ADT/SmallVector.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>

   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);

include/llvm/ADT/StringMap.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>

     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)...);
 

lib/IR/DataLayout.cpp

   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.

lib/Support/FoldingSet.cpp

 /// 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;

   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.

lib/Support/SmallPtrSet.cpp

 #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>

 
   // 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*));
 }
 

   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*));
 

   // 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.

         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);

lib/Support/SmallVector.cpp

   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;

lib/Support/StringMap.cpp

 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;

   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;