[CallGraph] Preserve call records vector when replacing call edge

Summary:
Try not to resize vector of call records in a call graph node when
replacing call edge. That would prevent invalidation of iterators
stored in the CG SCC pass manager's scc_iterator.

Reviewers: jdoerfert

Reviewed By: jdoerfert

Subscribers: hiraditya, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D84295

1 files changed, 29 insertions, 5 deletions

diff --git a/llvm/lib/Analysis/CallGraph.cpp b/llvm/lib/Analysis/CallGraph.cpp
index 55adb45..19c128b 100644
--- a/llvm/lib/Analysis/CallGraph.cpp
+++ b/llvm/lib/Analysis/CallGraph.cpp
@@ -281,13 +281,37 @@ void CallGraphNode::replaceCallEdge(CallBase &Call, CallBase &NewCall,
       I->second = NewNode;
       NewNode->AddRef();
 
-      // Refresh callback references.
-      forEachCallbackFunction(Call, [=](Function *CB) {
-        removeOneAbstractEdgeTo(CG->getOrInsertFunction(CB));
+      // Refresh callback references. Do not resize CalledFunctions if the
+      // number of callbacks is the same for new and old call sites.
+      SmallVector<CallGraphNode *, 4u> OldCBs;
+      SmallVector<CallGraphNode *, 4u> NewCBs;
+      forEachCallbackFunction(Call, [this, &OldCBs](Function *CB) {
+        OldCBs.push_back(CG->getOrInsertFunction(CB));
       });
-      forEachCallbackFunction(NewCall, [=](Function *CB) {
-        addCalledFunction(nullptr, CG->getOrInsertFunction(CB));
+      forEachCallbackFunction(NewCall, [this, &NewCBs](Function *CB) {
+        NewCBs.push_back(CG->getOrInsertFunction(CB));
       });
+      if (OldCBs.size() == NewCBs.size()) {
+        for (unsigned N = 0; N < OldCBs.size(); ++N) {
+          CallGraphNode *OldNode = OldCBs[N];
+          CallGraphNode *NewNode = NewCBs[N];
+          for (auto J = CalledFunctions.begin();; ++J) {
+            assert(J != CalledFunctions.end() &&
+                   "Cannot find callsite to update!");
+            if (!J->first && J->second == OldNode) {
+              J->second = NewNode;
+              OldNode->DropRef();
+              NewNode->AddRef();
+              break;
+            }
+          }
+        }
+      } else {
+        for (auto *CGN : OldCBs)
+          removeOneAbstractEdgeTo(CGN);
+        for (auto *CGN : NewCBs)
+          addCalledFunction(nullptr, CGN);
+      }
       return;
     }
   }