1 files changed, 9 insertions, 11 deletions

diff --git a/lib/Transforms/IPO/GlobalOpt.cpp b/lib/Transforms/IPO/GlobalOpt.cpp
index 3963fdea39..f11230c494 100644
--- a/lib/Transforms/IPO/GlobalOpt.cpp
+++ b/lib/Transforms/IPO/GlobalOpt.cpp
@@ -823,6 +823,7 @@ static void ConstantPropUsersOf(Value *V, LLVMContext &Context) {
 static GlobalVariable *OptimizeGlobalAddressOfMalloc(GlobalVariable *GV,
                                                      CallInst *CI,
                                                      BitCastInst *BCI,
+                                                     Value* NElems,
                                                      LLVMContext &Context,
                                                      TargetData* TD) {
   DEBUG(errs() << "PROMOTING MALLOC GLOBAL: " << *GV
@@ -830,9 +831,7 @@ static GlobalVariable *OptimizeGlobalAddressOfMalloc(GlobalVariable *GV,
 
   const Type *IntPtrTy = TD->getIntPtrType(Context);
   
-  Value* ArraySize = getMallocArraySize(CI, Context, TD);
-  assert(ArraySize && "not a malloc whose array size can be determined");
-  ConstantInt *NElements = cast<ConstantInt>(ArraySize);
+  ConstantInt *NElements = cast<ConstantInt>(NElems);
   if (NElements->getZExtValue() != 1) {
     // If we have an array allocation, transform it to a single element
     // allocation to make the code below simpler.
@@ -1275,15 +1274,14 @@ static void RewriteUsesOfLoadForHeapSRoA(LoadInst *Load,
 /// PerformHeapAllocSRoA - CI is an allocation of an array of structures.  Break
 /// it up into multiple allocations of arrays of the fields.
 static GlobalVariable *PerformHeapAllocSRoA(GlobalVariable *GV,
-                                            CallInst *CI, BitCastInst* BCI, 
+                                            CallInst *CI, BitCastInst* BCI,
+                                            Value* NElems,
                                             LLVMContext &Context,
-                                            TargetData *TD){
+                                            TargetData *TD) {
   DEBUG(errs() << "SROA HEAP ALLOC: " << *GV << "  MALLOC CALL = " << *CI 
                << " BITCAST = " << *BCI << '\n');
   const Type* MAT = getMallocAllocatedType(CI);
   const StructType *STy = cast<StructType>(MAT);
-  Value* ArraySize = getMallocArraySize(CI, Context, TD);
-  assert(ArraySize && "not a malloc whose array size can be determined");
 
   // There is guaranteed to be at least one use of the malloc (storing
   // it into GV).  If there are other uses, change them to be uses of
@@ -1309,7 +1307,7 @@ static GlobalVariable *PerformHeapAllocSRoA(GlobalVariable *GV,
     FieldGlobals.push_back(NGV);
     
     Value *NMI = CallInst::CreateMalloc(CI, TD->getIntPtrType(Context),
-                                        FieldTy, ArraySize,
+                                        FieldTy, NElems,
                                         BCI->getName() + ".f" + Twine(FieldNo));
     FieldMallocs.push_back(NMI);
     new StoreInst(NMI, NGV, BCI);
@@ -1510,7 +1508,7 @@ static bool TryToOptimizeStoreOfMallocToGlobal(GlobalVariable *GV,
       // something.
       if (TD && 
           NElements->getZExtValue() * TD->getTypeAllocSize(AllocTy) < 2048) {
-        GVI = OptimizeGlobalAddressOfMalloc(GV, CI, BCI, Context, TD);
+        GVI = OptimizeGlobalAddressOfMalloc(GV, CI, BCI, NElems, Context, TD);
         return true;
       }
   
@@ -1520,7 +1518,7 @@ static bool TryToOptimizeStoreOfMallocToGlobal(GlobalVariable *GV,
 
     // If this is an allocation of a fixed size array of structs, analyze as a
     // variable size array.  malloc [100 x struct],1 -> malloc struct, 100
-    if (!isArrayMalloc(CI, Context, TD))
+    if (NElems == ConstantInt::get(CI->getOperand(1)->getType(), 1))
       if (const ArrayType *AT = dyn_cast<ArrayType>(AllocTy))
         AllocTy = AT->getElementType();
   
@@ -1547,7 +1545,7 @@ static bool TryToOptimizeStoreOfMallocToGlobal(GlobalVariable *GV,
           CI = extractMallocCallFromBitCast(NewMI);
         }
       
-        GVI = PerformHeapAllocSRoA(GV, CI, BCI, Context, TD);
+        GVI = PerformHeapAllocSRoA(GV, CI, BCI, NElems, Context, TD);
         return true;
       }
     }