Teach instCombine to remove malloc+free if malloc's only uses are comparisons

to null.  Patch by Matti Niemenmaa.

llvm-svn: 104871

lib/Transforms/InstCombine/InstCombine.h

@@ -178,6 +178,7 @@ public:
   Instruction *visitPHINode(PHINode &PN);
   Instruction *visitGetElementPtrInst(GetElementPtrInst &GEP);
   Instruction *visitAllocaInst(AllocaInst &AI);
+  Instruction *visitMalloc(Instruction &FI);
   Instruction *visitFree(Instruction &FI);
   Instruction *visitLoadInst(LoadInst &LI);
   Instruction *visitStoreInst(StoreInst &SI);

lib/Transforms/InstCombine/InstCombineCalls.cpp

@@ -250,6 +250,8 @@ Instruction *InstCombiner::SimplifyMemSet(MemSetInst *MI) {
 Instruction *InstCombiner::visitCallInst(CallInst &CI) {
   if (isFreeCall(&CI))
     return visitFree(CI);
+  if (isMalloc(&CI))
+    return visitMalloc(CI);
 
   // If the caller function is nounwind, mark the call as nounwind, even if the
   // callee isn't.

lib/Transforms/InstCombine/InstCombineCompares.cpp

@@ -1924,35 +1924,6 @@ Instruction *InstCombiner::visitICmpInst(ICmpInst &I) {
         }
         break;
       }
-      case Instruction::Call:
-        // If we have (malloc != null), and if the malloc has a single use, we
-        // can assume it is successful and remove the malloc.
-        if (isMalloc(LHSI) && LHSI->hasOneUse() &&
-            isa<ConstantPointerNull>(RHSC)) {
-          // Need to explicitly erase malloc call here, instead of adding it to
-          // Worklist, because it won't get DCE'd from the Worklist since
-          // isInstructionTriviallyDead() returns false for function calls.
-          // It is OK to replace LHSI/MallocCall with Undef because the 
-          // instruction that uses it will be erased via Worklist.
-          if (extractMallocCall(LHSI)) {
-            LHSI->replaceAllUsesWith(UndefValue::get(LHSI->getType()));
-            EraseInstFromFunction(*LHSI);
-            return ReplaceInstUsesWith(I,
-                               ConstantInt::get(Type::getInt1Ty(I.getContext()),
-                                                      !I.isTrueWhenEqual()));
-          }
-          if (CallInst* MallocCall = extractMallocCallFromBitCast(LHSI))
-            if (MallocCall->hasOneUse()) {
-              MallocCall->replaceAllUsesWith(
-                                        UndefValue::get(MallocCall->getType()));
-              EraseInstFromFunction(*MallocCall);
-              Worklist.Add(LHSI); // The malloc's bitcast use.
-              return ReplaceInstUsesWith(I,
-                               ConstantInt::get(Type::getInt1Ty(I.getContext()),
-                                                      !I.isTrueWhenEqual()));
-            }
-        }
-        break;
       case Instruction::IntToPtr:
         // icmp pred inttoptr(X), null -> icmp pred X, 0
         if (RHSC->isNullValue() && TD &&

lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -710,6 +710,52 @@ Instruction *InstCombiner::visitGetElementPtrInst(GetElementPtrInst &GEP) {
   return 0;
 }
 
+
+
+static bool IsOnlyNullComparedAndFreed(const Value &V) {
+  for (Value::const_use_iterator UI = V.use_begin(), UE = V.use_end();
+       UI != UE; ++UI) {
+    if (isFreeCall(*UI))
+      continue;
+    if (const ICmpInst *ICI = dyn_cast<ICmpInst>(*UI))
+      if (ICI->isEquality() && isa<ConstantPointerNull>(ICI->getOperand(1)))
+        continue;
+    return false;
+  }
+  return true;
+}
+
+Instruction *InstCombiner::visitMalloc(Instruction &MI) {
+  // If we have a malloc call which is only used in any amount of comparisons
+  // to null and free calls, delete the calls and replace the comparisons with
+  // true or false as appropriate.
+  if (IsOnlyNullComparedAndFreed(MI)) {
+    for (Value::use_iterator UI = MI.use_begin(), UE = MI.use_end();
+         UI != UE;) {
+      // We can assume that every remaining use is a free call or an icmp eq/ne
+      // to null, so the cast is safe.
+      Instruction *I = cast<Instruction>(*UI);
+
+      // Early increment here, as we're about to get rid of the user.
+      ++UI;
+
+      if (isFreeCall(I)) {
+        EraseInstFromFunction(*cast<CallInst>(I));
+        continue;
+      }
+      // Again, the cast is safe.
+      ICmpInst *C = cast<ICmpInst>(I);
+      ReplaceInstUsesWith(*C, ConstantInt::get(Type::getInt1Ty(C->getContext()),
+                                               C->isFalseWhenEqual()));
+      EraseInstFromFunction(*C);
+    }
+    return EraseInstFromFunction(MI);
+  }
+  return 0;
+}
+
+
+
 Instruction *InstCombiner::visitFree(Instruction &FI) {
   Value *Op = FI.getOperand(1);
 
@@ -726,23 +772,6 @@ Instruction *InstCombiner::visitFree(Instruction &FI) {
   if (isa<ConstantPointerNull>(Op))
     return EraseInstFromFunction(FI);
 
-  // If we have a malloc call whose only use is a free call, delete both.
-  if (isMalloc(Op)) {
-    if (CallInst* CI = extractMallocCallFromBitCast(Op)) {
-      if (Op->hasOneUse() && CI->hasOneUse()) {
-        EraseInstFromFunction(FI);
-        EraseInstFromFunction(*CI);
-        return EraseInstFromFunction(*cast<Instruction>(Op));
-      }
-    } else {
-      // Op is a call to malloc
-      if (Op->hasOneUse()) {
-        EraseInstFromFunction(FI);
-        return EraseInstFromFunction(*cast<Instruction>(Op));
-      }
-    }
-  }
-
   return 0;
 }
 

test/Transforms/InstCombine/badmalloc.ll

@@ -10,6 +10,7 @@ declare void @free(i8*)
 define i1 @test1() {
   %A = call noalias i8* @malloc(i64 4) nounwind
   %B = icmp eq i8* %A, null
+  store i8 0, i8* %A
 
   call void @free(i8* %A)
   ret i1 %B

test/Transforms/InstCombine/malloc-free-delete.ll

@@ -1,13 +1,25 @@
-; RUN: opt < %s -instcombine -globaldce -S | FileCheck %s
+; RUN: opt < %s -instcombine -S | FileCheck %s
 ; PR1201
 define i32 @main(i32 %argc, i8** %argv) {
-        %c_19 = alloca i8*              ; <i8**> [#uses=2]
-        %malloc_206 = malloc i8, i32 10         ; <i8*> [#uses=1]
+        %c_19 = alloca i8*
+        %malloc_206 = malloc i8, i32 10
 ; CHECK-NOT: malloc
         store i8* %malloc_206, i8** %c_19
-        %tmp_207 = load i8** %c_19              ; <i8*> [#uses=1]
+        %tmp_207 = load i8** %c_19
         free i8* %tmp_207
 ; CHECK-NOT: free
         ret i32 0
 ; CHECK: ret i32 0
 }
+
+declare i8* @malloc(i32)
+declare void @free(i8*)
+
+define i1 @foo() {
+; CHECK: @foo
+; CHECK-NEXT: ret i1 false
+  %m = call i8* @malloc(i32 1)
+  %z = icmp eq i8* %m, null
+  call void @free(i8* %m)
+  ret i1 %z
+}