[clang] [llvm] [Coroutines] Never collect return value alloca into coroutine frame (PR #213580)

[via cfe-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-coroutines

@llvm/pr-subscribers-llvm-transforms

Author: Weibo He (NewSigma)

<details>
<summary>Changes</summary>

The coroutine return value must not reside within the coroutine frame; otherwise, a heap-use-after-free error will occur, as the frame is destroyed before the return is completed.

In the front end, emit `lifetime.start` for the return value alloca so that its lifetime begins after any suspension point. Do not emit the corresponding `lifetime.end`, because Clang considers `lifetime.end` just before the `ret` instruction is redundant.

In the middle end, remove the assumption that an address can cross suspension points if missing `lifetime.end` markers. This basically assumes that `ret` implies `lifetime.end` and aligns with front end's expectation.

Close #<!-- -->49843

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


5 Files Affected:

- (modified) clang/lib/CodeGen/CGCoroutine.cpp (+4) 
- (modified) clang/test/CodeGenCoroutines/coro-gro.cpp (+1) 
- (added) clang/test/CodeGenCoroutines/coro-gro5.cpp (+35) 
- (modified) llvm/lib/Transforms/Coroutines/SpillUtils.cpp (-5) 
- (modified) llvm/test/Transforms/Coroutines/coro-lifetime-end.ll (+2-2) 


``````````diff
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index bf896f1338ab4..e64778b36590c 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -880,6 +880,10 @@ struct GetReturnObjectManager {
     Builder.CreateCondBr(InRamp, ConvBB, AfterConvBB);
 
     CGF.EmitBlock(ConvBB);
+    // Emit lifetime.start after all suspending points, ensuring the return
+    // alloca does not go into the coroutine frame.
+    if (auto *AI = dyn_cast<llvm::AllocaInst>(CGF.ReturnValue.getBasePointer()))
+      CGF.EmitLifetimeStart(AI);
     CGF.EmitAnyExprToMem(S.getReturnValue(), CGF.ReturnValue,
                          S.getReturnValue()->getType().getQualifiers(),
                          /*IsInit*/ true);
diff --git a/clang/test/CodeGenCoroutines/coro-gro.cpp b/clang/test/CodeGenCoroutines/coro-gro.cpp
index e9f79e10db7ae..17065bb5ae729 100644
--- a/clang/test/CodeGenCoroutines/coro-gro.cpp
+++ b/clang/test/CodeGenCoroutines/coro-gro.cpp
@@ -61,6 +61,7 @@ int f() {
   // CHECK-NEXT: br i1 %InRamp, label %[[GroConv:.+]], label %[[AfterGroConv:.+]]
 
   // CHECK: [[GroConv]]:
+  // CHECK-NEXT: call void @llvm.lifetime.start.p0(ptr %[[RetVal]])
   // CHECK-NEXT: %[[Conv:.+]] = call noundef i32 @_ZN7GroTypecviEv(
   // CHECK-NEXT: store i32 %[[Conv]], ptr %[[RetVal]]
   // CHECK-NEXT: %[[IsActive:.+]] = load i1, ptr %[[GroActive]]
diff --git a/clang/test/CodeGenCoroutines/coro-gro5.cpp b/clang/test/CodeGenCoroutines/coro-gro5.cpp
new file mode 100644
index 0000000000000..a546d1e5aeade
--- /dev/null
+++ b/clang/test/CodeGenCoroutines/coro-gro5.cpp
@@ -0,0 +1,35 @@
+// Test that return value alloca does not enter the coro frame
+// Regression test for GH49843
+// RUN: %clang_cc1 -std=c++20 -triple=x86_64-unknown-linux-gnu -fexceptions -fcxx-exceptions -emit-llvm -o - %s | FileCheck %s
+
+#include "Inputs/coroutine.h"
+
+struct tag { char data[8]; }; // `tag` can be any type. It could be empty, or an int, or anything. 
+
+struct expected {
+  char data; // No issues if this member isn't here.
+
+  expected(tag) : data() {}
+
+  struct promise_type {
+    tag get_return_object() { return {}; } // No issues if we return an `expected` instead. 
+    std::suspend_never initial_suspend() { return {}; }
+    std::suspend_never final_suspend() noexcept { return {}; }
+    tag return_value(tag) { return tag{}; }
+    void unhandled_exception() {}
+  };
+};
+
+// CHECK-LABEL: define {{.*}} i8 @_Z2f1v()
+expected f1() {
+  // CHECK: %[[Retval:.+]] = alloca %struct.expected, align 1
+
+  // CHECK: gro.conv:
+  // CHECK-NEXT: call void @llvm.lifetime.start.p0(ptr %[[Retval]])
+  // CHECK: invoke void @_ZN8expectedC1E3tag(ptr {{.*}} %[[Retval]], i64 {{.*}})
+
+  // CHECK: %[[GEP:.+]] = getelementptr {{.*}} %struct.expected, ptr %[[Retval]], i32 0, i32 0
+  // CHECK-NEXT: %[[Val:.+]] = load i8, ptr %[[GEP]], align 1
+  // CHECK-NEXT: ret i8 %[[Val]]
+  co_return {};
+}
diff --git a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
index 05abccf0f9a97..8396430f5b4f0 100644
--- a/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
+++ b/llvm/lib/Transforms/Coroutines/SpillUtils.cpp
@@ -340,11 +340,6 @@ struct AllocaUseVisitor : PtrUseVisitor<AllocaUseVisitor> {
     // every basic block that uses the pointer to see if they cross suspension
     // points. The uses cover both direct uses as well as indirect uses.
     if (ShouldUseLifetimeStartInfo && !LifetimeStarts.empty()) {
-      // If there is no explicit lifetime.end, then assume the address can
-      // cross suspension points.
-      if (LifetimeEndBBs.empty())
-        return true;
-
       // If there is a path from a lifetime.start to a suspend without a
       // corresponding lifetime.end, then the alloca's lifetime persists
       // beyond that suspension point and the alloca must go on the frame.
diff --git a/llvm/test/Transforms/Coroutines/coro-lifetime-end.ll b/llvm/test/Transforms/Coroutines/coro-lifetime-end.ll
index ea044daa75533..9a9b40cc05a77 100644
--- a/llvm/test/Transforms/Coroutines/coro-lifetime-end.ll
+++ b/llvm/test/Transforms/Coroutines/coro-lifetime-end.ll
@@ -8,8 +8,8 @@ declare void @consume.i8.array(ptr)
 
 @testbool = external local_unnamed_addr global i8, align 1
 
-; testval does not contain an explicit lifetime end. We must assume that it may
-; live across suspension.
+; testval does not contain an explicit lifetime.end between lifetime.start and coro.suspend
+; We must assume that it may live across suspension.
 define void @HasNoLifetimeEnd() presplitcoroutine {
 ; CHECK-LABEL: define void @HasNoLifetimeEnd() {
 ; CHECK-NEXT:  entry:

``````````

</details>


https://github.com/llvm/llvm-project/pull/213580