1 files changed, 96 insertions, 112 deletions

diff --git a/lib/Analysis/MemoryBuiltins.cpp b/lib/Analysis/MemoryBuiltins.cpp
index cf2ad62b61..0982345cfa 100644
--- a/lib/Analysis/MemoryBuiltins.cpp
+++ b/lib/Analysis/MemoryBuiltins.cpp
@@ -88,64 +88,119 @@ const CallInst* llvm::extractMallocCallFromBitCast(const Value* I) {
                                       : NULL;
 }
 
-static bool isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
-                                const TargetData* TD) {
+/// isConstantOne - Return true only if val is constant int 1.
+static bool isConstantOne(Value *val) {
+  return isa<ConstantInt>(val) && cast<ConstantInt>(val)->isOne();
+}
+
+static Value* isArrayMallocHelper(const CallInst *CI, LLVMContext &Context,
+                                  const TargetData* TD) {
   if (!CI)
-    return false;
+    return NULL;
 
+  // Type must be known to determine array size.
   const Type* T = getMallocAllocatedType(CI);
-
-  // We can only indentify an array malloc if we know the type of the malloc 
-  // call.
-  if (!T) return false;
+  if (!T)
+    return NULL;
 
   Value* MallocArg = CI->getOperand(1);
-  Constant *ElementSize = ConstantExpr::getSizeOf(T);
+  ConstantExpr* CO = dyn_cast<ConstantExpr>(MallocArg);
+  BinaryOperator* BO = dyn_cast<BinaryOperator>(MallocArg);
+
+  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);
+  Constant *FoldedElementSize =
+   ConstantFoldConstantExpression(cast<ConstantExpr>(ElementSize), Context, TD);
+
+  // First, check if CI is a non-array malloc.
+  if (CO && ((CO == ElementSize) ||
+             (FoldedElementSize && (CO == FoldedElementSize))))
+    // Match CreateMalloc's use of constant 1 array-size for non-array mallocs.
+    return ConstantInt::get(MallocArg->getType(), 1);
+
+  // Second, check if CI is an array malloc whose array size can be determined.
+  if (isConstantOne(ElementSize) || 
+      (FoldedElementSize && isConstantOne(FoldedElementSize)))
+    return MallocArg;
 
-  BinaryOperator *BI = dyn_cast<BinaryOperator>(MallocArg);
-  if (!BI)
-    return false;
+  if (!CO && !BO)
+    return NULL;
 
-  if (BI->getOpcode() == Instruction::Mul)
-    // ArraySize * ElementSize
-    if (BI->getOperand(1) == ElementSize ||
-        (FoldedElementSize && BI->getOperand(1) == FoldedElementSize))
-      return true;
+  Value* Op0 = NULL;
+  Value* Op1 = NULL;
+  unsigned Opcode = 0;
+  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) || 
+             (BO->getOpcode() == Instruction::Shl))) {
+    Op0 = BO->getOperand(0);
+    Op1 = BO->getOperand(1);
+    Opcode = BO->getOpcode();
+  }
 
-  // TODO: Detect case where MallocArg mul has been transformed to shl.
+  // Determine array size if malloc's argument is the product of a mul or shl.
+  if (Op0) {
+    if (Opcode == Instruction::Mul) {
+      if ((Op1 == ElementSize) ||
+          (FoldedElementSize && (Op1 == FoldedElementSize)))
+        // ArraySize * ElementSize
+        return Op0;
+      if ((Op0 == ElementSize) ||
+          (FoldedElementSize && (Op0 == FoldedElementSize)))
+        // ElementSize * ArraySize
+        return Op1;
+    }
+    if (Opcode == Instruction::Shl) {
+      ConstantInt* Op1Int = dyn_cast<ConstantInt>(Op1);
+      if (!Op1Int) return NULL;
+      Value* Op1Pow = ConstantInt::get(Op1->getType(),
+                                       pow(2, Op1Int->getZExtValue()));
+      if (Op0 == ElementSize || (FoldedElementSize && Op0 == FoldedElementSize))
+        // ArraySize << log2(ElementSize)
+        return Op1Pow;
+      if (Op1Pow == ElementSize ||
+        (FoldedElementSize && Op1Pow == FoldedElementSize))
+        // ElementSize << log2(ArraySize)
+        return Op0;
+    }
+  }
 
-  return false;
+  // We could not determine the malloc array size from MallocArg.
+  return NULL;
 }
 
 /// isArrayMalloc - Returns the corresponding CallInst if the instruction 
-/// 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 size of the malloced type
-/// or
-///  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.
+/// is a call to malloc whose array size can be determined and the array size
+/// is not constant 1.  Otherwise, return NULL.
 CallInst* llvm::isArrayMalloc(Value* I, LLVMContext &Context,
                               const TargetData* TD) {
   CallInst *CI = extractMallocCall(I);
-  return (isArrayMallocHelper(CI, Context, TD)) ? CI : NULL;
+  Value* ArraySize = isArrayMallocHelper(CI, Context, TD);
+
+  if (ArraySize &&
+      ArraySize != ConstantInt::get(CI->getOperand(1)->getType(), 1))
+    return CI;
+
+  // CI is a non-array malloc or we can't figure out that it is an array malloc.
+  return NULL;
 }
 
 const CallInst* llvm::isArrayMalloc(const Value* I, LLVMContext &Context,
                                     const TargetData* TD) {
   const CallInst *CI = extractMallocCall(I);
-  return (isArrayMallocHelper(CI, Context, TD)) ? CI : NULL;
+  Value* ArraySize = isArrayMallocHelper(CI, Context, TD);
+
+  if (ArraySize &&
+      ArraySize != ConstantInt::get(CI->getOperand(1)->getType(), 1))
+    return CI;
+
+  // CI is a non-array malloc or we can't figure out that it is an array malloc.
+  return NULL;
 }
 
 /// getMallocType - Returns the PointerType resulting from the malloc call.
@@ -183,85 +238,14 @@ 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.  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.
-/// For non-array mallocs, the computed size is constant 1. 
-/// This function returns NULL for all mallocs whose array size cannot be
+/// getMallocArraySize - Returns the array size of a malloc call.  If the 
+/// argument passed to malloc is a multiple of the size of the malloced type,
+/// then return that multiple.  For non-array mallocs, the multiple is
+/// constant 1.  Otherwise, return NULL for 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);
-
-  Value* MallocArg = CI->getOperand(1);
-  assert(getMallocAllocatedType(CI) && "getMallocArraySize and no type");
-  Constant *ElementSize = ConstantExpr::getSizeOf(getMallocAllocatedType(CI));
-  ElementSize = ConstantExpr::getTruncOrBitCast(ElementSize, 
-                                                MallocArg->getType());
-
-  Constant* CO = dyn_cast<Constant>(MallocArg);
-  BinaryOperator* BO = dyn_cast<BinaryOperator>(MallocArg);
-  assert((isConstantOne(ElementSize) || CO || BO) &&
-         "getMallocArraySize and malformed malloc IR");
-      
-  if (isConstantOne(ElementSize))
-    return MallocArg;
-    
-  if (CO)
-    return CO->getOperand(0);
-    
-  // TODO: Detect case where MallocArg mul has been transformed to shl.
-
-  assert(BO && "getMallocArraySize not constant but not multiplication either");
-  return BO->getOperand(0);
+  return isArrayMallocHelper(CI, Context, TD);
 }
 
 //===----------------------------------------------------------------------===//