[Clang][NFC][Coroutine] Do not leave dangling pointers after removing CoroBegin

[yuxuanchen1997]

In handleNoSuspendCoroutine we delete the coro.begin intrinsic in lieu of splitting. After this operation Shape.CoroBegin becomes dangling. A new patch I am working on needs to tell that this happened.

After setting the dangling pointer to nullptr, I found that the only usage for the dangling pointer in the code is updateCallGraphAfterCoroutineSplit. After some investigation, the logic that checks CoroBegin's existence is not very useful. You would still need to remove CoroEnds and do the necessary updates regardless.

[llvmbot]

@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-llvm-transforms

Author: Yuxuan Chen (yuxuanchen1997)

Changes

In handleNoSuspendCoroutine we delete the coro.begin intrinsic in lieu of splitting. After this operation Shape.CoroBegin becomes dangling. A new patch I am working on needs to tell that this happened.

After setting the dangling pointer to nullptr, I found that the only usage for the dangling pointer in the code is updateCallGraphAfterCoroutineSplit. After some investigation, the logic that checks CoroBegin's existence is not very useful. You would still need to remove CoroEnds and do the necessary updates regardless.

---

Full diff: https://github.com/llvm/llvm-project/pull/98546.diff

1 Files Affected:

• (modified) llvm/lib/Transforms/Coroutines/CoroSplit.cpp (+13-15)

diff --git a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
index 0b52d1e4490cb..9e4da5f8ca961 100644
--- a/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
+++ b/llvm/lib/Transforms/Coroutines/CoroSplit.cpp
@@ -1248,6 +1248,7 @@ static void handleNoSuspendCoroutine(coro::Shape &Shape) {
   }
 
   CoroBegin->eraseFromParent();
+  Shape.CoroBegin = nullptr;
 }
 
 // SimplifySuspendPoint needs to check that there is no calls between
@@ -1970,9 +1971,17 @@ splitCoroutine(Function &F, SmallVectorImpl<Function *> &Clones,
 }
 
 /// Remove calls to llvm.coro.end in the original function.
-static void removeCoroEnds(const coro::Shape &Shape) {
-  for (auto *End : Shape.CoroEnds) {
-    replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
+static void removeCoroEndsFromRampFunction(const coro::Shape &Shape) {
+  if (Shape.ABI != coro::ABI::Switch) {
+    for (auto *End : Shape.CoroEnds) {
+      replaceCoroEnd(End, Shape, Shape.FramePtr, /*in resume*/ false, nullptr);
+    }
+  } else {
+    for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
+      auto &Context = End->getContext();
+      End->replaceAllUsesWith(ConstantInt::getFalse(Context));
+      End->eraseFromParent();
+    }
   }
 }
 
@@ -1981,18 +1990,6 @@ static void updateCallGraphAfterCoroutineSplit(
     const SmallVectorImpl<Function *> &Clones, LazyCallGraph::SCC &C,
     LazyCallGraph &CG, CGSCCAnalysisManager &AM, CGSCCUpdateResult &UR,
     FunctionAnalysisManager &FAM) {
-  if (!Shape.CoroBegin)
-    return;
-
-  if (Shape.ABI != coro::ABI::Switch)
-    removeCoroEnds(Shape);
-  else {
-    for (llvm::AnyCoroEndInst *End : Shape.CoroEnds) {
-      auto &Context = End->getContext();
-      End->replaceAllUsesWith(ConstantInt::getFalse(Context));
-      End->eraseFromParent();
-    }
-  }
 
   if (!Clones.empty()) {
     switch (Shape.ABI) {
@@ -2120,6 +2117,7 @@ PreservedAnalyses CoroSplitPass::run(LazyCallGraph::SCC &C,
     const coro::Shape Shape =
         splitCoroutine(F, Clones, FAM.getResult<TargetIRAnalysis>(F),
                        OptimizeFrame, MaterializableCallback);
+    removeCoroEndsFromRampFunction(Shape);
     updateCallGraphAfterCoroutineSplit(*N, Shape, Clones, C, CG, AM, UR, FAM);
 
     ORE.emit([&]() {

[ChuanqiXu9]

LGTM