diff --git a/llvm/include/llvm/Analysis/CallGraph.h b/llvm/include/llvm/Analysis/CallGraph.h --- a/llvm/include/llvm/Analysis/CallGraph.h +++ b/llvm/include/llvm/Analysis/CallGraph.h @@ -138,6 +138,10 @@ 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. diff --git a/llvm/lib/Analysis/CallGraph.cpp b/llvm/lib/Analysis/CallGraph.cpp --- a/llvm/lib/Analysis/CallGraph.cpp +++ b/llvm/lib/Analysis/CallGraph.cpp @@ -127,6 +127,16 @@ 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 diff --git a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp --- a/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp +++ b/llvm/lib/Transforms/Utils/CallGraphUpdater.cpp @@ -114,6 +114,7 @@ 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);