1 files changed, 44 insertions, 9 deletions
diff --git a/lib/Analysis/MallocHelper.cpp b/lib/Analysis/MallocHelper.cpp
index 89051d1788387..e7bb41eeec287 100644
--- a/lib/Analysis/MallocHelper.cpp
+++ b/lib/Analysis/MallocHelper.cpp
@@ -130,9 +130,9 @@ static bool isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
 /// matches the malloc call IR generated by CallInst::CreateMalloc().  This 
 /// means that it is a malloc call with one bitcast use AND the malloc call's 
 /// size argument is:
-///  1. a constant not equal to the malloc's allocated type
+///  1. a constant not equal to the size of the malloced type
 /// or
-///  2. the result of a multiplication by the malloc's allocated type
+///  2. the result of a multiplication by the size of the malloced type
 /// Otherwise it returns NULL.
 /// The unique bitcast is needed to determine the type/size of the array
 /// allocation.
@@ -183,25 +183,60 @@ const Type* llvm::getMallocAllocatedType(const CallInst* CI) {
   return PT ? PT->getElementType() : NULL;
 }
 
+/// isSafeToGetMallocArraySize - Returns true if the array size of a malloc can
+/// be determined.  It can be determined in these 3 cases of malloc codegen:
+/// 1. non-array malloc: The malloc's size argument is a constant and equals the ///    size of the type being malloced.
+/// 2. array malloc: This is a malloc call with one bitcast use AND the malloc
+///    call's size argument is a constant multiple of the size of the malloced
+///    type.
+/// 3. array malloc: This is a malloc call with one bitcast use AND the malloc
+///    call's size argument is the result of a multiplication by the size of the
+///    malloced type.
+/// Otherwise returns false.
+static bool isSafeToGetMallocArraySize(const CallInst *CI,
+                                       LLVMContext &Context,
+                                       const TargetData* TD) {
+  if (!CI)
+    return false;
+
+  // Type must be known to determine array size.
+  const Type* T = getMallocAllocatedType(CI);
+  if (!T) return false;
+
+  Value* MallocArg = CI->getOperand(1);
+  Constant *ElementSize = ConstantExpr::getSizeOf(T);
+  ElementSize = ConstantExpr::getTruncOrBitCast(ElementSize, 
+                                                MallocArg->getType());
+
+  // First, check if it is a non-array malloc.
+  if (isa<ConstantExpr>(MallocArg) && (MallocArg == ElementSize))
+    return true;
+
+  // Second, check if it can be determined that this is an array malloc.
+  return isArrayMallocHelper(CI, Context, TD);
+}
+
 /// isConstantOne - Return true only if val is constant int 1.
 static bool isConstantOne(Value *val) {
   return isa<ConstantInt>(val) && cast<ConstantInt>(val)->isOne();
 }
 
-/// getMallocArraySize - Returns the array size of a malloc call.  The array
-/// size is computated in 1 of 3 ways:
-///  1. If the element type if of size 1, then array size is the argument to 
+/// getMallocArraySize - Returns the array size of a malloc call.  For array
+/// mallocs, the size is computated in 1 of 3 ways:
+///  1. If the element type is of size 1, then array size is the argument to 
 ///     malloc.
 ///  2. Else if the malloc's argument is a constant, the array size is that
 ///     argument divided by the element type's size.
 ///  3. Else the malloc argument must be a multiplication and the array size is
 ///     the first operand of the multiplication.
-/// This function returns constant 1 if:
-///  1. The malloc call's allocated type cannot be determined.
-///  2. IR wasn't created by a call to CallInst::CreateMalloc() with a non-NULL
-///     ArraySize.
+/// For non-array mallocs, the computed size is constant 1. 
+/// This function returns NULL for all mallocs whose array size cannot be
+/// determined.
 Value* llvm::getMallocArraySize(CallInst* CI, LLVMContext &Context,
                                 const TargetData* TD) {
+  if (!isSafeToGetMallocArraySize(CI, Context, TD))
+    return NULL;
+
   // Match CreateMalloc's use of constant 1 array-size for non-array mallocs.
   if (!isArrayMalloc(CI, Context, TD))
     return ConstantInt::get(CI->getOperand(1)->getType(), 1);