1 files changed, 35 insertions, 41 deletions

diff --git a/lib/Analysis/MemoryBuiltins.cpp b/lib/Analysis/MemoryBuiltins.cpp
index f4eb793..e710350 100644
--- a/lib/Analysis/MemoryBuiltins.cpp
+++ b/lib/Analysis/MemoryBuiltins.cpp
@@ -17,7 +17,6 @@
 #include "llvm/Instructions.h"
 #include "llvm/Module.h"
 #include "llvm/Analysis/ConstantFolding.h"
-#include "llvm/Target/TargetData.h"
 using namespace llvm;
 
 //===----------------------------------------------------------------------===//
@@ -97,47 +96,45 @@ static Value *isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
   if (!CI)
     return NULL;
 
-  // The size of the malloc's result type must be known to determine array size.
+  // Type must be known to determine array size.
   const Type *T = getMallocAllocatedType(CI);
-  if (!T || !T->isSized() || !TD)
+  if (!T)
     return NULL;
 
   Value *MallocArg = CI->getOperand(1);
-  const Type *ArgType = MallocArg->getType();
   ConstantExpr *CO = dyn_cast<ConstantExpr>(MallocArg);
   BinaryOperator *BO = dyn_cast<BinaryOperator>(MallocArg);
 
-  unsigned ElementSizeInt = TD->getTypeAllocSize(T);
-  if (const StructType *ST = dyn_cast<StructType>(T))
-    ElementSizeInt = TD->getStructLayout(ST)->getSizeInBytes();
-  Constant *ElementSize = ConstantInt::get(ArgType, ElementSizeInt);
+  Constant *ElementSize = ConstantExpr::getSizeOf(T);
+  ElementSize = ConstantExpr::getTruncOrBitCast(ElementSize, 
+                                                MallocArg->getType());
+  Constant *FoldedElementSize =
+   ConstantFoldConstantExpression(cast<ConstantExpr>(ElementSize), Context, TD);
 
   // First, check if CI is a non-array malloc.
-  if (CO && CO == ElementSize)
+  if (CO && ((CO == ElementSize) ||
+             (FoldedElementSize && (CO == FoldedElementSize))))
     // Match CreateMalloc's use of constant 1 array-size for non-array mallocs.
-    return ConstantInt::get(ArgType, 1);
+    return ConstantInt::get(MallocArg->getType(), 1);
 
   // Second, check if CI is an array malloc whose array size can be determined.
-  if (isConstantOne(ElementSize))
+  if (isConstantOne(ElementSize) || 
+      (FoldedElementSize && isConstantOne(FoldedElementSize)))
     return MallocArg;
 
-  if (ConstantInt *CInt = dyn_cast<ConstantInt>(MallocArg))
-    if (CInt->getZExtValue() % ElementSizeInt == 0)
-      return ConstantInt::get(ArgType, CInt->getZExtValue() / ElementSizeInt);
-
   if (!CO && !BO)
     return NULL;
 
   Value *Op0 = NULL;
   Value *Op1 = NULL;
   unsigned Opcode = 0;
-  if (CO && ((CO->getOpcode() == Instruction::Mul) ||
+  if (CO && ((CO->getOpcode() == Instruction::Mul) || 
              (CO->getOpcode() == Instruction::Shl))) {
     Op0 = CO->getOperand(0);
     Op1 = CO->getOperand(1);
     Opcode = CO->getOpcode();
   }
-  if (BO && ((BO->getOpcode() == Instruction::Mul) ||
+  if (BO && ((BO->getOpcode() == Instruction::Mul) || 
              (BO->getOpcode() == Instruction::Shl))) {
     Op0 = BO->getOperand(0);
     Op1 = BO->getOperand(1);
@@ -147,10 +144,12 @@ static Value *isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
   // Determine array size if malloc's argument is the product of a mul or shl.
   if (Op0) {
     if (Opcode == Instruction::Mul) {
-      if (Op1 == ElementSize)
+      if ((Op1 == ElementSize) ||
+          (FoldedElementSize && (Op1 == FoldedElementSize)))
         // ArraySize * ElementSize
         return Op0;
-      if (Op0 == ElementSize)
+      if ((Op0 == ElementSize) ||
+          (FoldedElementSize && (Op0 == FoldedElementSize)))
         // ElementSize * ArraySize
         return Op1;
     }
@@ -162,10 +161,11 @@ static Value *isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
       uint64_t BitToSet = Op1Int.getLimitedValue(Op1Int.getBitWidth() - 1);
       Value *Op1Pow = ConstantInt::get(Context, 
                                   APInt(Op1Int.getBitWidth(), 0).set(BitToSet));
-      if (Op0 == ElementSize)
+      if (Op0 == ElementSize || (FoldedElementSize && Op0 == FoldedElementSize))
         // ArraySize << log2(ElementSize)
         return Op1Pow;
-      if (Op1Pow == ElementSize)
+      if (Op1Pow == ElementSize ||
+          (FoldedElementSize && Op1Pow == FoldedElementSize))
         // ElementSize << log2(ArraySize)
         return Op0;
     }
@@ -205,41 +205,35 @@ const CallInst *llvm::isArrayMalloc(const Value *I, LLVMContext &Context,
 }
 
 /// getMallocType - Returns the PointerType resulting from the malloc call.
-/// The PointerType depends on the number of bitcast uses of the malloc call:
-///   0: PointerType is the calls' return type.
-///   1: PointerType is the bitcast's result type.
-///  >1: Unique PointerType cannot be determined, return NULL.
+/// This PointerType is the result type of the call's only bitcast use.
+/// If there is no unique bitcast use, then return NULL.
 const PointerType *llvm::getMallocType(const CallInst *CI) {
   assert(isMalloc(CI) && "GetMallocType and not malloc call");
   
-  const PointerType *MallocType = NULL;
-  unsigned NumOfBitCastUses = 0;
-
+  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 (const BitCastInst *BCI = dyn_cast<BitCastInst>(*UI++)) {
-      MallocType = cast<PointerType>(BCI->getDestTy());
-      NumOfBitCastUses++;
-    }
+    if ((BCI = dyn_cast<BitCastInst>(cast<Instruction>(*UI++))))
+      break;
 
-  // Malloc call has 1 bitcast use, so type is the bitcast's destination type.
-  if (NumOfBitCastUses == 1)
-    return MallocType;
+  // 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());
 
   // Malloc call was not bitcast, so type is the malloc function's return type.
-  if (NumOfBitCastUses == 0)
+  if (!BCI)
     return cast<PointerType>(CI->getType());
 
   // Type could not be determined.
   return NULL;
 }
 
-/// getMallocAllocatedType - Returns the Type allocated by malloc call.
-/// The Type depends on the number of bitcast uses of the malloc call:
-///   0: PointerType is the malloc calls' return type.
-///   1: PointerType is the bitcast's result type.
-///  >1: Unique PointerType cannot be determined, return NULL.
+/// getMallocAllocatedType - Returns the Type allocated by malloc call. This
+/// Type is the result type of the call's only bitcast use. If there is no
+/// unique bitcast use, then return NULL.
 const Type *llvm::getMallocAllocatedType(const CallInst *CI) {
   const PointerType *PT = getMallocType(CI);
   return PT ? PT->getElementType() : NULL;