[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 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;
     }
   }
 }
 
 // Provide an explicit template instantiation for the static ID.
 AnalysisKey CallGraphAnalysis::Key;
 

llvm/unittests/IR/LegacyPassManagerTest.cpp

 // where possible.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/AbstractCallSite.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/OptBisect.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace llvm {
   void initializeModuleNDMPass(PassRegistry&);
TEST(PassManager, CallGraphUpdater0) {
       ASSERT_EQ(P->SetupWorked, 1U);
       ASSERT_EQ(P->NumSCCs, 4U);
       ASSERT_EQ(P->NumFns, 6U);
       ASSERT_EQ(P->NumFnDecls, 1U);
       ASSERT_EQ(M->getFunctionList().size(), 3U);
       ASSERT_EQ(P->NumExtCalledBefore, /* test1, 2a, 2b, 3, 4 */ 5U);
       ASSERT_EQ(P->NumExtCalledAfter, /* test1, 3repl, 4 */ 3U);
     }
+
+    // Test for call graph SCC pass that replaces all callback call instructions
+    // with clones and updates CallGraph by calling CallGraph::replaceCallEdge()
+    // method. Test is expected to complete successfully after running pass on
+    // all SCCs in the test module.
+    struct CallbackCallsModifierPass : public CGPass {
+      bool runOnSCC(CallGraphSCC &SCC) override {
+        CGPass::run();
+
+        CallGraph &CG = const_cast<CallGraph &>(SCC.getCallGraph());
+
+        bool Changed = false;
+        for (CallGraphNode *CGN : SCC) {
+          Function *F = CGN->getFunction();
+          if (!F || F->isDeclaration())
+            continue;
+
+          SmallVector<CallBase *, 4u> Calls;
+          for (Use &U : F->uses()) {
+            AbstractCallSite ACS(&U);
+            if (!ACS || !ACS.isCallbackCall() || !ACS.isCallee(&U))
+              continue;
+            Calls.push_back(cast<CallBase>(ACS.getInstruction()));
+          }
+          if (Calls.empty())
+            continue;
+
+          for (CallBase *OldCB : Calls) {
+            CallGraphNode *CallerCGN = CG[OldCB->getParent()->getParent()];
+            assert(any_of(*CallerCGN,
+                          [CGN](const CallGraphNode::CallRecord &CallRecord) {
+                            return CallRecord.second == CGN;
+                          }) &&
+                   "function is not a callee");
+
+            CallBase *NewCB = cast<CallBase>(OldCB->clone());
+
+            NewCB->insertBefore(OldCB);
+            NewCB->takeName(OldCB);
+
+            CallerCGN->replaceCallEdge(*OldCB, *NewCB, CG[F]);
+
+            OldCB->replaceAllUsesWith(NewCB);
+            OldCB->eraseFromParent();
+          }
+          Changed = true;
+        }
+        return Changed;
+      }
+    };
+
+    TEST(PassManager, CallbackCallsModifier0) {
+      LLVMContext Context;
+
+      const char *IR = "define void @foo() {\n"
+                       "  call void @broker(void (i8*)* @callback0, i8* null)\n"
+                       "  call void @broker(void (i8*)* @callback1, i8* null)\n"
+                       "  ret void\n"
+                       "}\n"
+                       "\n"
+                       "declare !callback !0 void @broker(void (i8*)*, i8*)\n"
+                       "\n"
+                       "define internal void @callback0(i8* %arg) {\n"
+                       "  ret void\n"
+                       "}\n"
+                       "\n"
+                       "define internal void @callback1(i8* %arg) {\n"
+                       "  ret void\n"
+                       "}\n"
+                       "\n"
+                       "!0 = !{!1}\n"
+                       "!1 = !{i64 0, i64 1, i1 false}";
+
+      SMDiagnostic Err;
+      std::unique_ptr<Module> M = parseAssemblyString(IR, Err, Context);
+      if (!M)
+        Err.print("LegacyPassManagerTest", errs());
+
+      CallbackCallsModifierPass *P = new CallbackCallsModifierPass();
+      legacy::PassManager Passes;
+      Passes.add(P);
+      Passes.run(*M);
+    }
   }
 }
 
 INITIALIZE_PASS(ModuleNDM, "mndm", "mndm", false, false)
 INITIALIZE_PASS_BEGIN(CGPass, "cgp","cgp", false, false)
 INITIALIZE_PASS_DEPENDENCY(CallGraphWrapperPass)
 INITIALIZE_PASS_END(CGPass, "cgp","cgp", false, false)
 INITIALIZE_PASS(FPass, "fp","fp", false, false)
 INITIALIZE_PASS_BEGIN(LPass, "lp","lp", false, false)
 INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_END(LPass, "lp","lp", false, false)