[Coroutines] properly update CallGraph in CoroSplit (#107935)

Fixes #107139.

We weren't updating the call graph properly in CoroSplit. This crash is
due to the await_suspend() function calling the coroutine, forming a
multi-node SCC. The issue bisected to
#79712 but I think this is red
herring. We haven't been properly updating the call graph.

Added an example of such code as a test case.

Author: yuxuanchen1997

llvm/lib/Transforms/Coroutines/CoroSplit.cpp

@@ -2080,12 +2080,13 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
   return Shape;
 }
 
-static void updateCallGraphAfterCoroutineSplit(
+static LazyCallGraph::SCC &updateCallGraphAfterCoroutineSplit(
     LazyCallGraph::Node &N, const coro::Shape &Shape,
     const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
     LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
     FunctionAnalysisManager &FAM) {
 
+  auto *CurrentSCC = &C;
   if (!Clones.empty()) {
     switch (Shape.ABI) {
     case coro::ABI::Switch:
@@ -2105,13 +2106,16 @@ static void updateCallGraphAfterCoroutineSplit(
     }
 
     // Let the CGSCC infra handle the changes to the original function.
-    updateCGAndAnalysisManagerForCGSCCPass(CG, C, N, AM, UR, FAM);
+    CurrentSCC = &updateCGAndAnalysisManagerForCGSCCPass(CG, *CurrentSCC, N, AM,
+                                                         UR, FAM);
   }
 
   // Do some cleanup and let the CGSCC infra see if we've cleaned up any edges
   // to the split functions.
   postSplitCleanup(N.getFunction());
-  updateCGAndAnalysisManagerForFunctionPass(CG, C, N, AM, UR, FAM);
+  CurrentSCC = &updateCGAndAnalysisManagerForFunctionPass(CG, *CurrentSCC, N,
+                                                          AM, UR, FAM);
+  return *CurrentSCC;
 }
 
 /// Replace a call to llvm.coro.prepare.retcon.
@@ -2200,6 +2204,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
   if (Coroutines.empty() && PrepareFns.empty())
     return PreservedAnalyses::all();
 
+  auto *CurrentSCC = &C;
   // Split all the coroutines.
   for (LazyCallGraph::Node *N : Coroutines) {
     Function &F = N->getFunction();
@@ -2211,7 +2216,8 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     coro::Shape Shape =
         splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
-    updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
+    CurrentSCC = &updateCallGraphAfterCoroutineSplit(
+        *N, Shape, Clones, *CurrentSCC, CG, AM, UR, FAM);
 
     auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
     ORE.emit([&]() {
@@ -2223,14 +2229,14 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
 
     if (!Shape.CoroSuspends.empty()) {
       // Run the CGSCC pipeline on the original and newly split functions.
-      UR.CWorklist.insert(&C);
+      UR.CWorklist.insert(CurrentSCC);
       for (Function *Clone : Clones)
         UR.CWorklist.insert(CG.lookupSCC(CG.get(*Clone)));
     }
   }
 
   for (auto *PrepareFn : PrepareFns) {
-    replaceAllPrepares(PrepareFn, CG, C);
+    replaceAllPrepares(PrepareFn, CG, *CurrentSCC);
   }
 
   return PreservedAnalyses::none();

llvm/test/Transforms/Coroutines/gh107139-split-in-scc.ll

@@ -0,0 +1,38 @@
+; Verify that we don't crash on mutually recursive coroutines
+; RUN: opt < %s -passes='cgscc(coro-split)' -S | FileCheck %s
+
+target triple = "x86_64-redhat-linux-gnu"
+
+; CHECK-LABEL: define void @foo
+define void @foo() presplitcoroutine personality ptr null {
+entry:
+
+  %0 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
+  %1 = call ptr @llvm.coro.begin(token %0, ptr null)
+  %2 = call token @llvm.coro.save(ptr null)
+  %3 = call i8 @llvm.coro.suspend(token none, i1 false)
+  %4 = call token @llvm.coro.save(ptr null)
+  ; CHECK: call void @bar(ptr null, ptr null)
+  call void @llvm.coro.await.suspend.void(ptr null, ptr null, ptr @bar)
+  ret void
+}
+
+; CHECK-LABEL: define void @bar({{.*}})
+define void @bar(ptr %0, ptr %1) {
+entry:
+  ; CHECK: call void @foo()
+  call void @foo()
+  ret void
+}
+
+; CHECK-LABEL: @foo.resume({{.*}})
+; CHECK-LABEL: @foo.destroy({{.*}})
+; CHECK-LABEL: @foo.cleanup({{.*}})
+
+declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr) #0
+declare ptr @llvm.coro.begin(token, ptr writeonly) nounwind
+declare token @llvm.coro.save(ptr) nomerge nounwind
+declare void @llvm.coro.await.suspend.void(ptr, ptr, ptr)
+declare i8 @llvm.coro.suspend(token, i1) nounwind
+
+attributes #0 = { nocallback nofree nosync nounwind willreturn memory(argmem: read) }