1 files changed, 4 insertions, 7 deletions
diff --git a/lib/Transforms/IPO/PruneEH.cpp b/lib/Transforms/IPO/PruneEH.cpp
index 8fb904b1ea..94d6331ddd 100644
--- a/lib/Transforms/IPO/PruneEH.cpp
+++ b/lib/Transforms/IPO/PruneEH.cpp
@@ -222,13 +222,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()));
   }