[CallGraphUpdater] Update the ExternalCallingNode for node replacements

[jdoerfert]

Summary

While it is uncommon that the ExternalCallingNode needs to be updated,
it can happen. It is uncommon because most functions listed as callees
have external linkage, modifying them is usually not allowed. That said,
there are also internal functions that have, or better had, their
"address taken" at construction time. We conservatively assume various
uses cause the address "to be taken". Furthermore, the user might have
become dead at some point. As a consequence, transformations, e.g., the
Attributor, might be able to replace a function that is listed
as callee of the ExternalCallingNode.

Since there is no function corresponding to the ExternalCallingNode, we
did just remove the node from the callee list if we replaced it (so
far). Now it would be preferable to replace it if needed and remove it
otherwise. However, removing the node has implications on the CGSCC
iteration. Locally, that caused some other nodes to be never visited
but it is for sure possible other (bad) side effects can occur. As it
seems conservatively safe to keep the new node in the callee list we
will do that for now.

[hfinkel]

Is it difficult to create a test case for this?

[jdoerfert]

Add test case

[hfinkel]

LGTM

llvm/include/llvm/Analysis/CallGraph.h

public:
   /// Returns the \c CallGraphNode which is used to represent
   /// undetermined calls into the callgraph.
   CallGraphNode *getExternalCallingNode() const { return ExternalCallingNode; }
 
   CallGraphNode *getCallsExternalNode() const {
     return CallsExternalNode.get();
   }
 
+  /// Old node has been deleted, and New is to be used in its place, update the
+  /// ExternalCallingNode.
+  void ReplaceExternalCallEdge(CallGraphNode *Old, CallGraphNode *New);
+
   //===---------------------------------------------------------------------
   // Functions to keep a call graph up to date with a function that has been
   // modified.
   //
 
   /// Unlink the function from this module, returning it.
   ///
   /// Because this removes the function from the module, the call graph node is

llvm/lib/Analysis/CallGraph.cpp

void CallGraph::print(raw_ostream &OS) const {
   for (CallGraphNode *CN : Nodes)
     CN->print(OS);
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void CallGraph::dump() const { print(dbgs()); }
 #endif
 
+void CallGraph::ReplaceExternalCallEdge(CallGraphNode *Old,
+                                        CallGraphNode *New) {
+  for (auto &CR : ExternalCallingNode->CalledFunctions)
+    if (CR.second == Old) {
+      CR.second->DropRef();
+      CR.second = New;
+      CR.second->AddRef();
+    }
+}
+
 // removeFunctionFromModule - Unlink the function from this module, returning
 // it.  Because this removes the function from the module, the call graph node
 // is destroyed.  This is only valid if the function does not call any other
 // functions (ie, there are no edges in it's CGN).  The easiest way to do this
 // is to dropAllReferences before calling this.
 //
 Function *CallGraph::removeFunctionFromModule(CallGraphNode *CGN) {
   assert(CGN->empty() && "Cannot remove function from call "

llvm/lib/Transforms/Utils/CallGraphUpdater.cpp

else
     DeadFunctions.push_back(&DeadFn);
 }
 
 void CallGraphUpdater::replaceFunctionWith(Function &OldFn, Function &NewFn) {
   OldFn.removeDeadConstantUsers();
   ReplacedFunctions.insert(&OldFn);
   if (CG) {
     // Update the call graph for the newly promoted function.
-    // CG->spliceFunction(&OldFn, &NewFn);
     CallGraphNode *OldCGN = (*CG)[&OldFn];
     CallGraphNode *NewCGN = CG->getOrInsertFunction(&NewFn);
     NewCGN->stealCalledFunctionsFrom(OldCGN);
+    CG->ReplaceExternalCallEdge(OldCGN, NewCGN);
 
     // And update the SCC we're iterating as well.
     CGSCC->ReplaceNode(OldCGN, NewCGN);
   } else if (LCG) {
     // Directly substitute the functions in the call graph.
     LazyCallGraph::Node &OldLCGN = LCG->get(OldFn);
     SCC->getOuterRefSCC().replaceNodeFunction(OldLCGN, NewFn);
   }

llvm/unittests/IR/LegacyPassManagerTest.cpp

Module *makeLLVMModule(LLVMContext &Context) {
         /*Linkage=*/GlobalValue::ExternalLinkage,
         /*Name=*/"test2", mod);
       func_test2->setCallingConv(CallingConv::C);
       AttributeList func_test2_PAL;
       func_test2->setAttributes(func_test2_PAL);
 
       Function* func_test3 = Function::Create(
         /*Type=*/FuncTy_0,
-        /*Linkage=*/GlobalValue::ExternalLinkage,
+        /*Linkage=*/GlobalValue::InternalLinkage,
         /*Name=*/"test3", mod);
       func_test3->setCallingConv(CallingConv::C);
       AttributeList func_test3_PAL;
       func_test3->setAttributes(func_test3_PAL);
 
       Function* func_test4 = Function::Create(
         /*Type=*/FuncTy_2,
         /*Linkage=*/GlobalValue::ExternalLinkage,
Module *makeLLVMModule(LLVMContext &Context) {
         BasicBlock *label_bb =
             BasicBlock::Create(Context, "bb", func_test4, nullptr);
         BasicBlock *label_bb1 =
             BasicBlock::Create(Context, "bb1", func_test4, nullptr);
         BasicBlock *label_return =
             BasicBlock::Create(Context, "return", func_test4, nullptr);
 
         // Block entry (label_entry_11)
+        auto *AI = new AllocaInst(func_test3->getType(), 0, "func3ptr",
+                                  label_entry_11);
+        new StoreInst(func_test3, AI, label_entry_11);
         BranchInst::Create(label_bb, label_entry_11);
 
         // Block bb (label_bb)
         BranchInst::Create(label_bb, label_bb1, int1_f, label_bb);
 
         // Block bb1 (label_bb1)
         BranchInst::Create(label_bb1, label_return, int1_f, label_bb1);
 
Function *splitSimpleFunction(Function &F) {
       CI.setCalledFunction(SF);
       return SF;
     }
 
     struct CGModifierPass : public CGPass {
       unsigned NumSCCs = 0;
       unsigned NumFns = 0;
       bool SetupWorked = true;
+      unsigned NumExtCalledBefore = 0;
+      unsigned NumExtCalledAfter = 0;
 
       CallGraphUpdater CGU;
 
       bool runOnSCC(CallGraphSCC &SCMM) override {
         ++NumSCCs;
         for (CallGraphNode *N : SCMM)
           if (N->getFunction())
             ++NumFns;
 
         CGPass::run();
 
+        CallGraph &CG = const_cast<CallGraph &>(SCMM.getCallGraph());
+        CallGraphNode *ExtCallingNode = CG.getExternalCallingNode();
+        NumExtCalledBefore = ExtCallingNode->size();
+
         if (SCMM.size() <= 1)
           return false;
 
         CallGraphNode *N = *(SCMM.begin());
         Function *F = N->getFunction();
         Module *M = F->getParent();
         Function *Test1F = M->getFunction("test1");
         Function *Test2aF = M->getFunction("test2a");
         Function *Test2bF = M->getFunction("test2b");
         Function *Test3F = M->getFunction("test3");
+
         auto InSCC = [&](Function *Fn) {
           return llvm::any_of(SCMM, [Fn](CallGraphNode *CGN) {
             return CGN->getFunction() == Fn;
           });
         };
 
         if (!Test1F || !Test2aF || !Test2bF || !Test3F || !InSCC(Test1F) ||
             !InSCC(Test2aF) || !InSCC(Test2bF) || !InSCC(Test3F))
           return SetupWorked = false;
 
         CallInst *CI = dyn_cast<CallInst>(&Test1F->getEntryBlock().front());
         if (!CI || CI->getCalledFunction() != Test2aF)
           return SetupWorked = false;
 
-        CI->setCalledFunction(Test3F);
+        // Create a replica of test3 and just move the blocks there.
+        Function *Test3FRepl = Function::Create(
+            /*Type=*/Test3F->getFunctionType(),
+            /*Linkage=*/GlobalValue::InternalLinkage,
+            /*Name=*/"test3repl", Test3F->getParent());
+        while (!Test3F->empty()) {
+          BasicBlock &BB = Test3F->front();
+          BB.removeFromParent();
+          BB.insertInto(Test3FRepl);
+        }
+
+        CGU.initialize(CG, SCMM);
+
+        // Replace test3 with the replica. This is legal as it is actually
+        // internal and the "capturing use" is not really capturing anything.
+        CGU.replaceFunctionWith(*Test3F, *Test3FRepl);
+        Test3F->replaceAllUsesWith(Test3FRepl);
+
+        // Rewrite the call in test1 to point to the replica of 3 not test2.
+        CI->setCalledFunction(Test3FRepl);
 
-        CGU.initialize(const_cast<CallGraph &>(SCMM.getCallGraph()), SCMM);
+        // Delete test2a and test2b and reanalyze 1 as we changed calls inside.
         CGU.removeFunction(*Test2aF);
         CGU.removeFunction(*Test2bF);
         CGU.reanalyzeFunction(*Test1F);
+
         return true;
       }
 
-      bool doFinalization(CallGraph &CG) override { return CGU.finalize(); }
+      bool doFinalization(CallGraph &CG) override {
+        CGU.finalize();
+        // We removed test2 and replaced the internal test3.
+        NumExtCalledAfter = CG.getExternalCallingNode()->size();
+        return true;
+      }
     };
 
     TEST(PassManager, CallGraphUpdater0) {
       // SCC#1: test1->test2a->test2b->test3->test1
       // SCC#2: test4
       // SCC#3: indirect call node
 
       LLVMContext Context;
       std::unique_ptr<Module> M(makeLLVMModule(Context));
       ASSERT_EQ(M->getFunctionList().size(), 4U);
       Function *F = M->getFunction("test2");
       Function *SF = splitSimpleFunction(*F);
       CallInst::Create(F, "", &SF->getEntryBlock());
       ASSERT_EQ(M->getFunctionList().size(), 5U);
       CGModifierPass *P = new CGModifierPass();
       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(M->getFunctionList().size(), 3U);
+      ASSERT_EQ(P->NumExtCalledBefore, /* test1, 2a, 2b, 3, 4 */ 5U);
+      ASSERT_EQ(P->NumExtCalledAfter, /* test1, 3repl, 4 */ 3U);
     }
   }
 }
 
 INITIALIZE_PASS(ModuleNDM, "mndm", "mndm", false, false)
 INITIALIZE_PASS_BEGIN(CGPass, "cgp","cgp", false, false)
 INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
 INITIALIZE_PASS_END(CGPass, "cgp","cgp", false, false)
 INITIALIZE_PASS(FPass, "fp","fp", false, false)
 INITIALIZE_PASS_BEGIN(LPass, "lp","lp", false, false)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_END(LPass, "lp","lp", false, false)