[llvm] [CoroSplit] Fix use-after-free related to coro.suspend (PR #156572)

[Weibo He via llvm-commits]

https://github.com/NewSigma updated https://github.com/llvm/llvm-project/pull/156572

>From 6961a22d5d34f19ff4bca4fd6f6b8e7dcd637094 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Wed, 3 Sep 2025 10:17:02 +0800
Subject: [PATCH 1/2] [CoroSplit] Fix use after free of coro.suspend

---
 llvm/lib/Transforms/Coroutines/Coroutines.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Coroutines/Coroutines.cpp b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
index ac93f748ce65c..28a89a8f87dbd 100644
--- a/llvm/lib/Transforms/Coroutines/Coroutines.cpp
+++ b/llvm/lib/Transforms/Coroutines/Coroutines.cpp
@@ -356,9 +356,9 @@ void coro::Shape::invalidateCoroutine(
     // present.
     for (AnyCoroSuspendInst *CS : CoroSuspends) {
       CS->replaceAllUsesWith(PoisonValue::get(CS->getType()));
-      CS->eraseFromParent();
       if (auto *CoroSave = CS->getCoroSave())
         CoroSave->eraseFromParent();
+      CS->eraseFromParent();
     }
     CoroSuspends.clear();
 

>From 2fa1c66c3cc3d87576c3d877a768cfe184d63239 Mon Sep 17 00:00:00 2001
From: NewSigma <NewSigma at 163.com>
Date: Wed, 3 Sep 2025 11:03:25 +0800
Subject: [PATCH 2/2] Add test

---
 .../Transforms/Coroutines/coro-split-invalid.ll    | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 llvm/test/Transforms/Coroutines/coro-split-invalid.ll

diff --git a/llvm/test/Transforms/Coroutines/coro-split-invalid.ll b/llvm/test/Transforms/Coroutines/coro-split-invalid.ll
new file mode 100644
index 0000000000000..94fe539697214
--- /dev/null
+++ b/llvm/test/Transforms/Coroutines/coro-split-invalid.ll
@@ -0,0 +1,14 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Tests that coro-split correctly invalidate bad coroutines
+; RUN: opt < %s -passes='cgscc(coro-split)' -S | FileCheck %s
+
+define void @pr156444() presplitcoroutine {
+; CHECK-LABEL: define void @pr156444(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = call i8 @llvm.coro.suspend(token none, i1 false)
+  ret void
+}