[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.
diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index a8d7eb3..ef52775 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -5727,11 +5727,12 @@
     }
   }
 
+  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 @@
     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 @@
       }
 
       // 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 @@
           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 @@
     }
 
     // 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();
+    }
 
-    assert(OldFn->getNumUses() == 0 && "Unexpected leftover uses!");
-    OldFn->eraseFromParent();
+    ToBeDeletedFunctions.insert(OldFn);
+
     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 0000000..b81d354
--- /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
+}
+