Support, IR, ADT: Check nullptr after allocation with malloc/realloc or calloc

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

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@308576 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 40 insertions, 10 deletions

diff --git a/lib/IR/DataLayout.cpp b/lib/IR/DataLayout.cpp
index 5de281a9523..f4dddeb30d0 100644
--- a/lib/IR/DataLayout.cpp
+++ b/lib/IR/DataLayout.cpp
@@ -572,6 +572,8 @@ const StructLayout *DataLayout::getStructLayout(StructType *Ty) const {
   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.
diff --git a/lib/Support/FoldingSet.cpp b/lib/Support/FoldingSet.cpp
index 4496d06a15f..94237954903 100644
--- a/lib/Support/FoldingSet.cpp
+++ b/lib/Support/FoldingSet.cpp
@@ -215,6 +215,10 @@ static void **GetBucketFor(unsigned Hash, void **Buckets, unsigned NumBuckets) {
 /// 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;
@@ -271,10 +275,11 @@ void FoldingSetBase::GrowBucketCount(unsigned NewBucketCount) {
   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.
diff --git a/lib/Support/SmallPtrSet.cpp b/lib/Support/SmallPtrSet.cpp
index aa12e85fa4c..47e960e243c 100644
--- a/lib/Support/SmallPtrSet.cpp
+++ b/lib/Support/SmallPtrSet.cpp
@@ -15,6 +15,7 @@
 #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>
@@ -32,7 +33,9 @@ void SmallPtrSetImplBase::shrink_and_clear() {
 
   // 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*));
 }
 
@@ -96,8 +99,12 @@ void SmallPtrSetImplBase::Grow(unsigned NewSize) {
   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*));
 
@@ -125,7 +132,8 @@ SmallPtrSetImplBase::SmallPtrSetImplBase(const void **SmallStorage,
   // 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.
@@ -162,7 +170,8 @@ void SmallPtrSetImplBase::CopyFrom(const SmallPtrSetImplBase &RHS) {
         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);
diff --git a/lib/Support/SmallVector.cpp b/lib/Support/SmallVector.cpp
index b931505bd6a..74313151c76 100644
--- a/lib/Support/SmallVector.cpp
+++ b/lib/Support/SmallVector.cpp
@@ -26,14 +26,17 @@ void SmallVectorBase::grow_pod(void *FirstEl, size_t 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;
diff --git a/lib/Support/StringMap.cpp b/lib/Support/StringMap.cpp
index d2315966e32..4341da2d97b 100644
--- a/lib/Support/StringMap.cpp
+++ b/lib/Support/StringMap.cpp
@@ -52,14 +52,21 @@ StringMapImpl::StringMapImpl(unsigned InitSize, unsigned itemSize) {
 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;
@@ -215,6 +222,10 @@ unsigned StringMapImpl::RehashTable(unsigned BucketNo) {
   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;