[CallGraphUpdater] Remove nodes from their SCC (old PM)

Summary:
We can and should remove deleted nodes from their respective SCCs. We
did not do this before and this was a potential problem even though I
couldn't locally trigger an issue. Since the `DeleteNode` would assert
if the node was not in the SCC, we know we only remove nodes from their
SCC and only once (when run on all the Attributor tests).

Reviewers: lebedev.ri, hfinkel, fhahn, probinson, wristow, loladiro, sstefan1, uenoku

Subscribers: hiraditya, bollu, uenoku, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D77855

include/llvm/Analysis/CallGraphSCCPass.h

@@ -103,6 +103,10 @@ public:
   /// Old node has been deleted, and New is to be used in its place.
   void ReplaceNode(CallGraphNode *Old, CallGraphNode *New);
 
+  /// DeleteNode - This informs the SCC and the pass manager that the specified
+  /// Old node has been deleted.
+  void DeleteNode(CallGraphNode *Old);
+
   using iterator = std::vector<CallGraphNode *>::const_iterator;
 
   iterator begin() const { return Nodes.begin(); }

lib/Analysis/CallGraphSCCPass.cpp

@@ -562,6 +562,10 @@ void CallGraphSCC::ReplaceNode(CallGraphNode *Old, CallGraphNode *New) {
   CGI->ReplaceNode(Old, New);
 }
 
+void CallGraphSCC::DeleteNode(CallGraphNode *Old) {
+  ReplaceNode(Old, /* New */ nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // CallGraphSCCPass Implementation
 //===----------------------------------------------------------------------===//

lib/Transforms/Utils/CallGraphUpdater.cpp

@@ -110,6 +110,10 @@ void CallGraphUpdater::removeFunction(Function &DeadFn) {
     DeadFunctionsInComdats.push_back(&DeadFn);
   else
     DeadFunctions.push_back(&DeadFn);
+
+  // For the old call graph we remove the function from the SCC right away.
+  if (CGSCC && !ReplacedFunctions.count(&DeadFn))
+    CGSCC->DeleteNode((*CG)[&DeadFn]);
 }
 
 void CallGraphUpdater::replaceFunctionWith(Function &OldFn, Function &NewFn) {

unittests/IR/LegacyPassManagerTest.cpp

@@ -581,7 +581,8 @@ namespace llvm {
     struct CGModifierPass : public CGPass {
       unsigned NumSCCs = 0;
       unsigned NumFns = 0;
-      bool SetupWorked = true;
+      unsigned NumFnDecls = 0;
+      unsigned SetupWorked = 0;
       unsigned NumExtCalledBefore = 0;
       unsigned NumExtCalledAfter = 0;
 
@@ -589,10 +590,12 @@ namespace llvm {
 
       bool runOnSCC(CallGraphSCC &SCMM) override {
         ++NumSCCs;
-        for (CallGraphNode *N : SCMM)
-          if (N->getFunction())
+        for (CallGraphNode *N : SCMM) {
+          if (N->getFunction()){
             ++NumFns;
-
+            NumFnDecls += N->getFunction()->isDeclaration();
+          }
+        }
         CGPass::run();
 
         CallGraph &CG = const_cast<CallGraph &>(SCMM.getCallGraph());
@@ -618,11 +621,13 @@ namespace llvm {
 
         if (!Test1F || !Test2aF || !Test2bF || !Test3F || !InSCC(Test1F) ||
             !InSCC(Test2aF) || !InSCC(Test2bF) || !InSCC(Test3F))
-          return SetupWorked = false;
+          return false;
 
         CallInst *CI = dyn_cast<CallInst>(&Test1F->getEntryBlock().front());
         if (!CI || CI->getCalledFunction() != Test2aF)
-          return SetupWorked = false;
+          return false;
+
+        SetupWorked += 1;
 
         // Create a replica of test3 and just move the blocks there.
         Function *Test3FRepl = Function::Create(
@@ -664,7 +669,11 @@ namespace llvm {
     TEST(PassManager, CallGraphUpdater0) {
       // SCC#1: test1->test2a->test2b->test3->test1
       // SCC#2: test4
-      // SCC#3: indirect call node
+      // SCC#3: test3 (the empty function declaration as we replaced it with
+      //               test3repl when we visited SCC#1)
+      // SCC#4: test2a->test2b (the empty function declarations as we deleted
+      //                        these functions when we visited SCC#1)
+      // SCC#5: indirect call node
 
       LLVMContext Context;
       std::unique_ptr<Module> M(makeLLVMModule(Context));
@@ -677,9 +686,10 @@ namespace llvm {
       legacy::PassManager Passes;
       Passes.add(P);
       Passes.run(*M);
-      ASSERT_TRUE(P->SetupWorked);
-      ASSERT_EQ(P->NumSCCs, 3U);
-      ASSERT_EQ(P->NumFns, 5U);
+      ASSERT_EQ(P->SetupWorked, 1U);
+      ASSERT_EQ(P->NumSCCs, 5U);
+      ASSERT_EQ(P->NumFns, 8U);
+      ASSERT_EQ(P->NumFnDecls, 3U);
       ASSERT_EQ(M->getFunctionList().size(), 3U);
       ASSERT_EQ(P->NumExtCalledBefore, /* test1, 2a, 2b, 3, 4 */ 5U);
       ASSERT_EQ(P->NumExtCalledAfter, /* test1, 3repl, 4 */ 3U);