[Attributor][FIX] Allow dead users of rewritten function

If we replace a function with a new one because we rewrite the
signature, dead users may still refer to the old version. With this
patch we reuse the code that deals with dead functions, which the old
versions are, to avoid problems.

2 files changed, 70 insertions, 13 deletions

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index a8d7eb35b98..ef52775c759 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -5727,11 +5727,12 @@ ChangeStatus Attributor::run(Module &M) {
     }
   }
 
+  STATS_DECL(AAIsDead, Function, "Number of dead functions deleted.");
+  BUILD_STAT_NAME(AAIsDead, Function) += ToBeDeletedFunctions.size();
+
   // Rewrite the functions as requested during manifest.
   ManifestChange = ManifestChange | rewriteFunctionSignatures();
 
-  STATS_DECL(AAIsDead, Function, "Number of dead functions deleted.");
-  BUILD_STAT_NAME(AAIsDead, Function) += ToBeDeletedFunctions.size();
   for (Function *Fn : ToBeDeletedFunctions) {
     Fn->deleteBody();
     Fn->replaceAllUsesWith(UndefValue::get(Fn->getType()));
@@ -5890,8 +5891,9 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
     NewFn->getBasicBlockList().splice(NewFn->begin(),
                                       OldFn->getBasicBlockList());
 
-    // Set of all "call-like" instructions that invoke the old function.
-    SmallPtrSet<Instruction *, 8> OldCallSites;
+    // Set of all "call-like" instructions that invoke the old function mapped
+    // to their new replacements.
+    SmallVector<std::pair<CallBase *, CallBase *>, 8> CallSitePairs;
 
     // Callback to create a new "call-like" instruction for a given one.
     auto CallSiteReplacementCreator = [&](AbstractCallSite ACS) {
@@ -5943,7 +5945,6 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
       }
 
       // Copy over various properties and the new attributes.
-      OldCB->replaceAllUsesWith(NewCB);
       uint64_t W;
       if (OldCB->extractProfTotalWeight(W))
         NewCB->setProfWeight(W);
@@ -5954,10 +5955,7 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
           Ctx, OldCallAttributeList.getFnAttributes(),
           OldCallAttributeList.getRetAttributes(), NewArgOperandAttributes));
 
-      bool Inserted = OldCallSites.insert(OldCB).second;
-      assert(Inserted && "Call site was old twice!");
-      (void)Inserted;
-
+      CallSitePairs.push_back({OldCB, NewCB});
       return true;
     };
 
@@ -5984,11 +5982,15 @@ ChangeStatus Attributor::rewriteFunctionSignatures() {
     }
 
     // Eliminate the instructions *after* we visited all of them.
-    for (Instruction *OldCallSite : OldCallSites)
-      OldCallSite->eraseFromParent();
+    for (auto &CallSitePair : CallSitePairs) {
+      CallBase &OldCB = *CallSitePair.first;
+      CallBase &NewCB = *CallSitePair.second;
+      OldCB.replaceAllUsesWith(&NewCB);
+      OldCB.eraseFromParent();
+    }
+
+    ToBeDeletedFunctions.insert(OldFn);
 
-    assert(OldFn->getNumUses() == 0 && "Unexpected leftover uses!");
-    OldFn->eraseFromParent();
     Changed = ChangeStatus::CHANGED;
   }
 
diff --git a/llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll b/llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll
new file mode 100644
index 00000000000..b81d35491f9
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/ArgumentPromotion/live_called_from_dead.ll
@@ -0,0 +1,55 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -basicaa -attributor -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=3 < %s | FileCheck %s --check-prefixes=CHECK,OLDPM_MODULE
+; RUN: opt -S -passes='attributor' -aa-pipeline='basic-aa' -attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=3 < %s | FileCheck %s --check-prefixes=CHECK,NEWPM_MODULE
+
+; OLDPM_MODULE-NOT: @dead
+; NEWPM_MODULE-NOT: @dead
+
+define internal void @dead() {
+  call i32 @test(i32* null, i32* null)
+  ret void
+}
+
+define internal i32 @test(i32* %X, i32* %Y) {
+; CHECK-LABEL: define {{[^@]+}}@test()
+; CHECK-NEXT:    br i1 true, label [[LIVE:%.*]], label [[DEAD:%.*]]
+; CHECK:       live:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       dead:
+; CHECK-NEXT:    unreachable
+;
+  br i1 true, label %live, label %dead
+live:
+  ret i32 0
+dead:
+  call i32 @caller(i32* null)
+  call void @dead()
+  ret i32 1
+}
+
+define internal i32 @caller(i32* %B) {
+; CHECK-LABEL: define {{[^@]+}}@caller()
+; CHECK-NEXT:    [[A:%.*]] = alloca i32
+; CHECK-NEXT:    store i32 1, i32* [[A]], align 4
+; CHECK-NEXT:    [[C:%.*]] = call i32 @test()
+; CHECK-NEXT:    ret i32 0
+;
+  %A = alloca i32
+  store i32 1, i32* %A
+  %C = call i32 @test(i32* %A, i32* %B)
+  ret i32 %C
+}
+
+define i32 @callercaller() {
+; CHECK-LABEL: define {{[^@]+}}@callercaller()
+; CHECK-NEXT:    [[B:%.*]] = alloca i32
+; CHECK-NEXT:    store i32 2, i32* [[B]], align 4
+; CHECK-NEXT:    [[X:%.*]] = call i32 @caller()
+; CHECK-NEXT:    ret i32 0
+;
+  %B = alloca i32
+  store i32 2, i32* %B
+  %X = call i32 @caller(i32* %B)
+  ret i32 %X
+}
+