1 files changed, 55 insertions, 41 deletions

diff --git a/lib/Analysis/MallocHelper.cpp b/lib/Analysis/MallocHelper.cpp
index 60930d5169..9c1671d306 100644
--- a/lib/Analysis/MallocHelper.cpp
+++ b/lib/Analysis/MallocHelper.cpp
@@ -16,6 +16,7 @@
 #include "llvm/Constants.h"
 #include "llvm/Instructions.h"
 #include "llvm/Module.h"
+#include "llvm/Analysis/ConstantFolding.h"
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
@@ -24,10 +25,6 @@ using namespace llvm;
 
 /// isMalloc - Returns true if the the value is either a malloc call or a
 /// bitcast of the result of a malloc call.
-bool llvm::isMalloc(Value* I) {
-  return extractMallocCall(I) || extractMallocCallFromBitCast(I);
-}
-
 bool llvm::isMalloc(const Value* I) {
   return extractMallocCall(I) || extractMallocCallFromBitCast(I);
 }
@@ -79,40 +76,42 @@ const CallInst* llvm::extractMallocCallFromBitCast(const Value* I) {
                                       : NULL;
 }
 
-static bool isArrayMallocHelper(const CallInst *CI) {
+static bool isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
+                                const TargetData* TD) {
   if (!CI)
     return false;
 
-  // Only identify array mallocs for mallocs with 1 bitcast use.  The unique 
-  // bitcast is needed to determine the type/size of the array allocation.
-  if (!CI->hasOneUse()) return false;
+  const Type* T = getMallocAllocatedType(CI);
 
-  for (Value::use_const_iterator UI = CI->use_begin(), E = CI->use_end();
-       UI != E; )
-    if (!isa<BitCastInst>(cast<Instruction>(*UI++)))
-      return false;
+  // We can only indentify an array malloc if we know the type of the malloc 
+  // call.
+  if (!T) return false;
 
-  // malloc arg
   Value* MallocArg = CI->getOperand(1);
-  // element size
-  const Type* T = getMallocAllocatedType(CI);
-  if (!T) return false;
   Constant *ElementSize = ConstantExpr::getSizeOf(T);
-  
+  ElementSize = ConstantExpr::getTruncOrBitCast(ElementSize, 
+                                                MallocArg->getType());
+  Constant *FoldedElementSize = ConstantFoldConstantExpression(
+                                       cast<ConstantExpr>(ElementSize), 
+                                       Context, TD);
+
+
   if (isa<ConstantExpr>(MallocArg))
-    return (MallocArg == ElementSize) ? false : true;
+    return (MallocArg != ElementSize);
 
   BinaryOperator *BI = dyn_cast<BinaryOperator>(MallocArg);
   if (!BI)
     return false;
 
-  if (BI->getOpcode() != Instruction::Mul)
-    return false;
-      
-  if (BI->getOperand(1) != ElementSize)
-    return false;
-        
-  return true;
+  if (BI->getOpcode() == Instruction::Mul)
+    // ArraySize * ElementSize
+    if (BI->getOperand(1) == ElementSize ||
+        (FoldedElementSize && BI->getOperand(1) == FoldedElementSize))
+      return true;
+
+  // TODO: Detect case where MallocArg mul has been transformed to shl.
+
+  return false;
 }
 
 /// isArrayMalloc - Returns the corresponding CallInst if the instruction 
@@ -125,14 +124,16 @@ static bool isArrayMallocHelper(const CallInst *CI) {
 /// Otherwise it returns NULL.
 /// The unique bitcast is needed to determine the type/size of the array
 /// allocation.
-CallInst* llvm::isArrayMalloc(Value* I) {
+CallInst* llvm::isArrayMalloc(Value* I, LLVMContext &Context,
+                              const TargetData* TD) {
   CallInst *CI = extractMallocCall(I);
-  return (isArrayMallocHelper(CI)) ? CI : NULL;
+  return (isArrayMallocHelper(CI, Context, TD)) ? CI : NULL;
 }
 
-const CallInst* llvm::isArrayMalloc(const Value* I) {
+const CallInst* llvm::isArrayMalloc(const Value* I, LLVMContext &Context,
+                                    const TargetData* TD) {
   const CallInst *CI = extractMallocCall(I);
-  return (isArrayMallocHelper(CI)) ? CI : NULL;
+  return (isArrayMallocHelper(CI, Context, TD)) ? CI : NULL;
 }
 
 /// getMallocType - Returns the PointerType resulting from the malloc call.
@@ -142,14 +143,24 @@ const PointerType* llvm::getMallocType(const CallInst* CI) {
   assert(isMalloc(CI) && "GetMallocType and not malloc call");
   
   const BitCastInst* BCI = NULL;
+  
+  // Determine if CallInst has a bitcast use.
+  for (Value::use_const_iterator UI = CI->use_begin(), E = CI->use_end();
+       UI != E; )
+    if ((BCI = dyn_cast<BitCastInst>(cast<Instruction>(*UI++))))
+      break;
+
+  // Malloc call has 1 bitcast use and no other uses, so type is the bitcast's
+  // destination type.
+  if (BCI && CI->hasOneUse())
+    return cast<PointerType>(BCI->getDestTy());
 
-  // Determine type only if there is only 1 bitcast use of CI.
-  if (CI->hasOneUse())
-    for (Value::use_const_iterator UI = CI->use_begin(), E = CI->use_end();
-         UI != E; )
-      BCI = dyn_cast<BitCastInst>(cast<Instruction>(*UI++));
+  // Malloc call was not bitcast, so the type is the malloc's return type, i8*.
+  if (!BCI)
+    return cast<PointerType>(CI->getType());
 
-  return BCI ? reinterpret_cast<const PointerType*>(BCI->getDestTy()) : NULL;
+  // Type could not be determined.
+  return NULL;
 }
 
 /// getMallocAllocatedType - Returns the Type allocated by malloc call. This
@@ -177,15 +188,16 @@ static bool isConstantOne(Value *val) {
 ///  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.
-Value* llvm::getMallocArraySize(CallInst* CI) {
+Value* llvm::getMallocArraySize(CallInst* CI, LLVMContext &Context,
+                                const TargetData* TD) {
   // Match CreateMalloc's use of constant 1 array-size for non-array mallocs.
-  if (!isArrayMalloc(CI))
+  if (!isArrayMalloc(CI, Context, TD))
     return ConstantInt::get(CI->getOperand(1)->getType(), 1);
 
   Value* MallocArg = CI->getOperand(1);
   assert(getMallocAllocatedType(CI) && "getMallocArraySize and no type");
   Constant *ElementSize = ConstantExpr::getSizeOf(getMallocAllocatedType(CI));
-  ElementSize = ConstantExpr::getTruncOrBitCast(cast<Constant>(ElementSize), 
+  ElementSize = ConstantExpr::getTruncOrBitCast(ElementSize, 
                                                 MallocArg->getType());
 
   Constant* CO = dyn_cast<Constant>(MallocArg);
@@ -195,10 +207,12 @@ Value* llvm::getMallocArraySize(CallInst* CI) {
       
   if (isConstantOne(ElementSize))
     return MallocArg;
+    
+  if (CO)
+    return CO->getOperand(0);
+    
+  // TODO: Detect case where MallocArg mul has been transformed to shl.
 
-  if (CO) 
-    return ConstantExpr::getUDiv(CO, ElementSize);
-  
   assert(BO && "getMallocArraySize not constant but not multiplication either");
   return BO->getOperand(0);
 }