[CallGraph] Preserve call records vector when replacing call edge

[sdmitriev]

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.

[jdoerfert]

This is a "fix" for the failure we see only if the number of callbacks is the same, correct?

[sdmitriev]

Yes, exactly. If the number of callbacks is the same we can avoid invalidating iterators. I assume that should really cover most of the cases.

[sdmitriev]

Added unit test for the CallGraphNode::replaceCallEdge() change. Without this change test finishes with an assertion on windows debug build.

[jdoerfert]

LGTM. This is not "great" but one of those shortcomings of the legacy CG we can work around in the hops it will be removed soon.
FWIW, in the absence of callbacks this should not cause much cost and be a no-op.

llvm/lib/Analysis/CallGraph.cpp

void CallGraphNode::replaceCallEdge(CallBase &Call, CallBase &NewCall,
   for (CalledFunctionsVector::iterator I = CalledFunctions.begin(); ; ++I) {
     assert(I != CalledFunctions.end() && "Cannot find callsite to remove!");
     if (I->first && *I->first == &Call) {
       I->second->DropRef();
       I->first = &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 I = 0; I < OldCBs.size(); I++) {
+          CallGraphNode *OldNode = OldCBs[I];
+          CallGraphNode *NewNode = NewCBs[I];
+          for (auto I = CalledFunctions.begin();; ++I) {
+            assert(I != CalledFunctions.end() &&
+                   "Cannot find callsite to update!");
+            if (!I->first && I->second == OldNode) {
+              I->second = NewNode;
+              OldNode->DropRef();
+              NewNode->AddRef();
+              break;
+            }
+          }
+        }
+      } else {
+        for (auto *CGN : OldCBs)
+          removeOneAbstractEdgeTo(CGN);
+        for (auto *CGN : NewCBs)
+          addCalledFunction(nullptr, CGN);
+      }
       return;
     }
   }
 }
 
 // Provide an explicit template instantiation for the static ID.
 AnalysisKey CallGraphAnalysis::Key;