1 files changed, 9 insertions, 0 deletions
diff --git a/lib/Analysis/PointerTracking.cpp b/lib/Analysis/PointerTracking.cpp
index e098647887..22818369b8 100644
--- a/lib/Analysis/PointerTracking.cpp
+++ b/lib/Analysis/PointerTracking.cpp
@@ -13,6 +13,7 @@
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/Dominators.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/MallocHelper.h"
 #include "llvm/Analysis/PointerTracking.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
@@ -99,6 +100,14 @@ const SCEV *PointerTracking::computeAllocationCount(Value *P,
     return SE->getSCEV(arraySize);
   }
 
+  if (CallInst *CI = extractMallocCall(V)) {
+    Value *arraySize = getMallocArraySize(CI, P->getContext(), TD);
+    Ty = getMallocAllocatedType(CI);
+    if (!Ty || !arraySize) return SE->getCouldNotCompute();
+    // arraySize elements of type Ty.
+    return SE->getSCEV(arraySize);
+  }
+
   if (GlobalVariable *GV = dyn_cast<GlobalVariable>(V)) {
     if (GV->hasDefinitiveInitializer()) {
       Constant *C = GV->getInitializer();