[Attributor][FIX] Do not apply h2s for arbitrary mallocs

H2S did apply to mallocs of non-constant sizes if the uses were OK. This
is now forbidden through reording of the "good" and "bad" cases in the
conditional.

llvm-svn: 374698

2 files changed, 27 insertions, 16 deletions

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index eb6756c53e7..6db3f40b05f 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -3620,30 +3620,36 @@ ChangeStatus AAHeapToStackImpl::updateImpl(Attributor &A) {
   };
 
   auto MallocCallocCheck = [&](Instruction &I) {
-    if (isMallocLikeFn(&I, TLI)) {
+    if (BadMallocCalls.count(&I))
+      return true;
+
+    bool IsMalloc = isMallocLikeFn(&I, TLI);
+    bool IsCalloc = !IsMalloc && isCallocLikeFn(&I, TLI);
+    if (!IsMalloc && !IsCalloc) {
+      BadMallocCalls.insert(&I);
+      return true;
+    }
+
+    if (IsMalloc) {
       if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(0)))
-        if (!Size->getValue().sle(MaxHeapToStackSize))
-          return true;
-    } else if (isCallocLikeFn(&I, TLI)) {
+        if (Size->getValue().sle(MaxHeapToStackSize))
+          if (UsesCheck(I)) {
+            MallocCalls.insert(&I);
+            return true;
+          }
+    } else if (IsCalloc) {
       bool Overflow = false;
       if (auto *Num = dyn_cast<ConstantInt>(I.getOperand(0)))
         if (auto *Size = dyn_cast<ConstantInt>(I.getOperand(1)))
-          if (!(Size->getValue().umul_ov(Num->getValue(), Overflow))
+          if ((Size->getValue().umul_ov(Num->getValue(), Overflow))
                    .sle(MaxHeapToStackSize))
-            if (!Overflow)
+            if (!Overflow && UsesCheck(I)) {
+              MallocCalls.insert(&I);
               return true;
-    } else {
-      BadMallocCalls.insert(&I);
-      return true;
+            }
     }
 
-    if (BadMallocCalls.count(&I))
-      return true;
-
-    if (UsesCheck(I))
-      MallocCalls.insert(&I);
-    else
-      BadMallocCalls.insert(&I);
+    BadMallocCalls.insert(&I);
     return true;
   };
 
diff --git a/llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll b/llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll
index 26a778be841..542b233770c 100644
--- a/llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll
+++ b/llvm/test/Transforms/FunctionAttrs/heap_to_stack.ll
@@ -316,3 +316,8 @@ define void @test14() {
   ; CHECK: tail call void @free(i8* noalias %1)
   ret void
 }
+
+define void @test15(i64 %S) {
+  %1 = tail call noalias i8* @malloc(i64 %S)
+  ret void
+}