1 files changed, 6 insertions, 0 deletions
diff --git a/lib/Transforms/IPO/PruneEH.cpp b/lib/Transforms/IPO/PruneEH.cpp
index 8fb904b1ea..630768fbdc 100644
--- a/lib/Transforms/IPO/PruneEH.cpp
+++ b/lib/Transforms/IPO/PruneEH.cpp
@@ -156,6 +156,9 @@ bool PruneEH::runOnSCC(const std::vector<CallGraphNode *> &SCC) {
 // function if we have invokes to non-unwinding functions or code after calls to
 // no-return functions.
 bool PruneEH::SimplifyFunction(Function *F) {
+  CallGraph &CG = getAnalysis<CallGraph>();
+  CallGraphNode *CGN = CG[F];
+
   bool MadeChange = false;
   for (Function::iterator BB = F->begin(), E = F->end(); BB != E; ++BB) {
     if (InvokeInst *II = dyn_cast<InvokeInst>(BB->getTerminator()))
@@ -174,6 +177,9 @@ bool PruneEH::SimplifyFunction(Function *F) {
         BasicBlock *UnwindBlock = II->getUnwindDest();
         UnwindBlock->removePredecessor(II->getParent());
 
+        // Fix up the call graph.
+        CGN->replaceCallSite(II, Call);
+
         // Insert a branch to the normal destination right before the
         // invoke.
         BranchInst::Create(II->getNormalDest(), II);