Reapply 55859. This doesn't change anything as

long as the callgraph is correct.  It checks
for wrong callgraphs more strictly.


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@55894 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 11 insertions, 13 deletions

diff --git a/lib/Transforms/IPO/PruneEH.cpp b/lib/Transforms/IPO/PruneEH.cpp
index 630768fbdc..297c6c1541 100644
--- a/lib/Transforms/IPO/PruneEH.cpp
+++ b/lib/Transforms/IPO/PruneEH.cpp
@@ -228,13 +228,10 @@ void PruneEH::DeleteBasicBlock(BasicBlock *BB) {
   CallGraphNode *CGN = CG[BB->getParent()];
   for (BasicBlock::iterator I = BB->end(), E = BB->begin(); I != E; ) {
     --I;
-    if (CallInst *CI = dyn_cast<CallInst>(I)) {
-      if (Function *Callee = CI->getCalledFunction())
-        CGN->removeCallEdgeTo(CG[Callee]);
-    } else if (InvokeInst *II = dyn_cast<InvokeInst>(I)) {
-      if (Function *Callee = II->getCalledFunction())
-        CGN->removeCallEdgeTo(CG[Callee]);
-    }
+    if (CallInst *CI = dyn_cast<CallInst>(I))
+      CGN->removeCallEdgeFor(CI);
+    else if (InvokeInst *II = dyn_cast<InvokeInst>(I))
+      CGN->removeCallEdgeFor(II);
     if (!I->use_empty())
       I->replaceAllUsesWith(UndefValue::get(I->getType()));
   }
diff --git a/lib/Transforms/Utils/InlineFunction.cpp b/lib/Transforms/Utils/InlineFunction.cpp
index cb9c28dcab..b0e70bb89a 100644
--- a/lib/Transforms/Utils/InlineFunction.cpp
+++ b/lib/Transforms/Utils/InlineFunction.cpp
@@ -143,16 +143,18 @@ static void HandleInlinedInvoke(InvokeInst *II, BasicBlock *FirstNewBlock,
 /// UpdateCallGraphAfterInlining - Once we have cloned code over from a callee
 /// into the caller, update the specified callgraph to reflect the changes we
 /// made.  Note that it's possible that not all code was copied over, so only
-/// some edges of the callgraph will be remain.
-static void UpdateCallGraphAfterInlining(const Function *Caller,
-                                         const Function *Callee,
+/// some edges of the callgraph may remain.
+static void UpdateCallGraphAfterInlining(CallSite CS,
                                          Function::iterator FirstNewBlock,
                                        DenseMap<const Value*, Value*> &ValueMap,
                                          CallGraph &CG) {
+  const Function *Caller = CS.getInstruction()->getParent()->getParent();
+  const Function *Callee = CS.getCalledFunction();
+
   // Update the call graph by deleting the edge from Callee to Caller
   CallGraphNode *CalleeNode = CG[Callee];
   CallGraphNode *CallerNode = CG[Caller];
-  CallerNode->removeCallEdgeTo(CalleeNode);
+  CallerNode->removeCallEdgeFor(CS);
 
   // Since we inlined some uninlined call sites in the callee into the caller,
   // add edges from the caller to all of the callees of the callee.
@@ -302,8 +304,7 @@ bool llvm::InlineFunction(CallSite CS, CallGraph *CG, const TargetData *TD) {
 
     // Update the callgraph if requested.
     if (CG)
-      UpdateCallGraphAfterInlining(Caller, CalledFunc, FirstNewBlock, ValueMap,
-                                   *CG);
+      UpdateCallGraphAfterInlining(CS, FirstNewBlock, ValueMap, *CG);
   }
 
   // If there are any alloca instructions in the block that used to be the entry