From a51c39cc3265f5d0d5de87b4a3ef9332c83556e1 Mon Sep 17 00:00:00 2001
From: Chris Lattner <sabre@nondot.org>
Date: Tue, 15 Sep 2009 05:40:35 +0000
Subject: [PATCH] add a new CallGraphNode::replaceCallEdge method and use it
 from argpromote to avoid invalidating an iterator.  This fixes PR4977. All
 clang tests now pass with expensive checking (on my system at least).

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@81843 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/llvm/Analysis/CallGraph.h        |  6 ++++++
 lib/Analysis/IPA/CallGraph.cpp           | 17 +++++++++++++++++
 lib/Transforms/IPO/ArgumentPromotion.cpp |  8 ++++----
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/include/llvm/Analysis/CallGraph.h b/include/llvm/Analysis/CallGraph.h
index 1d23c695ad4..bcb6dee033d 100644
--- a/include/llvm/Analysis/CallGraph.h
+++ b/include/llvm/Analysis/CallGraph.h
@@ -270,6 +270,12 @@ public:
   /// removeOneAbstractEdgeTo - Remove one edge associated with a null callsite
   /// from this node to the specified callee function.
   void removeOneAbstractEdgeTo(CallGraphNode *Callee);
+  
+  /// replaceCallEdge - This method replaces the edge in the node for the
+  /// specified call site with a new one.  Note that this method takes linear
+  /// time, so it should be used sparingly.
+  void replaceCallEdge(CallSite CS, CallSite NewCS, CallGraphNode *NewNode);
+  
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/lib/Analysis/IPA/CallGraph.cpp b/lib/Analysis/IPA/CallGraph.cpp
index 645916e35e7..e2b288d1ba9 100644
--- a/lib/Analysis/IPA/CallGraph.cpp
+++ b/lib/Analysis/IPA/CallGraph.cpp
@@ -279,5 +279,22 @@ void CallGraphNode::removeOneAbstractEdgeTo(CallGraphNode *Callee) {
   }
 }
 
+/// replaceCallEdge - This method replaces the edge in the node for the
+/// specified call site with a new one.  Note that this method takes linear
+/// time, so it should be used sparingly.
+void CallGraphNode::replaceCallEdge(CallSite CS,
+                                    CallSite NewCS, CallGraphNode *NewNode){
+  for (CalledFunctionsVector::iterator I = CalledFunctions.begin(); ; ++I) {
+    assert(I != CalledFunctions.end() && "Cannot find callsite to remove!");
+    if (I->first == CS.getInstruction()) {
+      I->second->DropRef();
+      I->first = NewCS.getInstruction();
+      I->second = NewNode;
+      NewNode->AddRef();
+      return;
+    }
+  }
+}
+
 // Enuse that users of CallGraph.h also link with this file
 DEFINING_FILE_FOR(CallGraph)
diff --git a/lib/Transforms/IPO/ArgumentPromotion.cpp b/lib/Transforms/IPO/ArgumentPromotion.cpp
index 3c584c894af..5b91f3d2099 100644
--- a/lib/Transforms/IPO/ArgumentPromotion.cpp
+++ b/lib/Transforms/IPO/ArgumentPromotion.cpp
@@ -589,7 +589,7 @@ CallGraphNode *ArgPromotion::DoPromotion(Function *F,
   // Construct the new function type using the new arguments.
   FunctionType *NFTy = FunctionType::get(RetTy, Params, FTy->isVarArg());
 
-  // Create the new function body and insert it into the module...
+  // Create the new function body and insert it into the module.
   Function *NF = Function::Create(NFTy, F->getLinkage(), F->getName());
   NF->copyAttributesFrom(F);
 
@@ -599,7 +599,8 @@ CallGraphNode *ArgPromotion::DoPromotion(Function *F,
   
   // Recompute the parameter attributes list based on the new arguments for
   // the function.
-  NF->setAttributes(AttrListPtr::get(AttributesVec.begin(), AttributesVec.end()));
+  NF->setAttributes(AttrListPtr::get(AttributesVec.begin(),
+                                     AttributesVec.end()));
   AttributesVec.clear();
 
   F->getParent()->getFunctionList().insert(F, NF);
@@ -729,8 +730,7 @@ CallGraphNode *ArgPromotion::DoPromotion(Function *F,
 
     // Update the callgraph to know that the callsite has been transformed.
     CallGraphNode *CalleeNode = CG[Call->getParent()->getParent()];
-    CalleeNode->removeCallEdgeFor(Call);
-    CalleeNode->addCalledFunction(New, NF_CGN);
+    CalleeNode->replaceCallEdge(Call, New, NF_CGN);
 
     if (!Call->use_empty()) {
       Call->replaceAllUsesWith(New);
-- 
2.34.1