[clang][CodeGen] Extends lifetime of coroutine promise(CWG2563)

NewSigma · NewSigma · commit 1999c901e563 · 2025-10-13T09:59:00.000+08:00
diff --git a/clang/docs/ReleaseNotes.rst b/clang/docs/ReleaseNotes.rst
@@ -169,6 +169,9 @@ C++20 Feature Support
   As a result, Clang no longer incorrectly diagnoses substitution failures in template arguments only
   used in concept-ids, and produces better diagnostics for satisfaction failure. (#GH61811) (#GH135190)
 
+- Fixed premature coroutine promise destruction when the coroutine completes
+  without suspending, as part of CWG2563. (#GH120200)
+
 C++17 Feature Support
 ^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -16,6 +16,7 @@
 #include "clang/AST/StmtCXX.h"
 #include "clang/AST/StmtVisitor.h"
 #include "llvm/ADT/ScopeExit.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 
 using namespace clang;
 using namespace CodeGen;
@@ -73,6 +74,13 @@ struct clang::CodeGen::CGCoroData {
   // the address of the coroutine frame of the current coroutine.
   llvm::CallInst *CoroBegin = nullptr;
 
+  // Stores the cloned cleanup block.
+  BasicBlock *RampCleanupBB = nullptr;
+
+  // Stores the llvm.coro.end that identifies if the coroutine exit without
+  // suspend.
+  llvm::CallInst *InRamp = nullptr;
+
   // Stores the last emitted coro.free for the deallocate expressions, we use it
   // to wrap dealloc code with if(auto mem = coro.free) dealloc(mem).
   llvm::CallInst *LastCoroFree = nullptr;
@@ -597,7 +605,7 @@ struct CallCoroEnd final : public EHScopeStack::Cleanup {
 namespace {
 // Make sure to call coro.delete on scope exit.
 struct CallCoroDelete final : public EHScopeStack::Cleanup {
-  Stmt *Deallocate;
+  const CoroutineBodyStmt *S;
 
   // Emit "if (coro.free(CoroId, CoroBegin)) Deallocate;"
 
@@ -607,16 +615,54 @@ struct CallCoroDelete final : public EHScopeStack::Cleanup {
   // builds a single call to a deallocation function which is safe to emit
   // multiple times.
   void Emit(CodeGenFunction &CGF, Flags) override {
+    bool Splited = CGF.CurCoro.Data->InRamp != nullptr;
+    if (!Splited)
+      splitCleanupBB(CGF);
+    EmitCoroFree(CGF);
+  }
+
+  // Clone the cleanup block for ramp function if the coroutine completes
+  // without suspending
+  void splitCleanupBB(CodeGenFunction &CGF) {
+    auto &Builder = CGF.Builder;
+    auto &CoroData = *CGF.CurCoro.Data;
+    auto *SaveInsertPt = Builder.GetInsertBlock();
+    auto *PreCleanupBB = SaveInsertPt->getSinglePredecessor();
+
+    auto *CleanupBB =
+        PreCleanupBB->splitBasicBlock(PreCleanupBB->begin(), "coro.cleanup");
+
+    PreCleanupBB->getTerminator()->eraseFromParent();
+    Builder.SetInsertPoint(PreCleanupBB);
+
+    llvm::Function *CoroInRamp =
+        CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_is_in_ramp);
+    CoroData.InRamp = Builder.CreateCall(CoroInRamp, {}, "InRamp");
+
+    BasicBlock *EndBB = CoroData.CleanupJD.getBlock();
+    Builder.CreateCondBr(CoroData.InRamp, EndBB, CleanupBB);
+    Builder.SetInsertPoint(SaveInsertPt);
+
+    if (S->getReturnStmt()) {
+      // Clone cleanup block before EmitCoroFree()
+      llvm::ValueToValueMapTy VMap{};
+      CoroData.RampCleanupBB =
+          llvm::CloneBasicBlock(CleanupBB, VMap, ".ramp");
+    }
+  }
+
+  void EmitCoroFree(CodeGenFunction &CGF, const Twine &NameSuffix = "") {
     // Remember the current point, as we are going to emit deallocation code
     // first to get to coro.free instruction that is an argument to a delete
     // call.
     BasicBlock *SaveInsertBlock = CGF.Builder.GetInsertBlock();
 
-    auto *FreeBB = CGF.createBasicBlock("coro.free");
+    Stmt *Deallocate = S->getDeallocate();
+    auto *FreeBB = CGF.createBasicBlock("coro.free" + NameSuffix);
     CGF.EmitBlock(FreeBB);
     CGF.EmitStmt(Deallocate);
 
-    auto *AfterFreeBB = CGF.createBasicBlock("after.coro.free");
+    auto *AfterFreeBB = CGF.createBasicBlock("after.coro.free" + NameSuffix);
     CGF.EmitBlock(AfterFreeBB);
 
     // We should have captured coro.free from the emission of deallocate.
@@ -641,7 +687,7 @@ struct CallCoroDelete final : public EHScopeStack::Cleanup {
     InsertPt->eraseFromParent();
     CGF.Builder.SetInsertPoint(AfterFreeBB);
   }
-  explicit CallCoroDelete(Stmt *DeallocStmt) : Deallocate(DeallocStmt) {}
+  explicit CallCoroDelete(const CoroutineBodyStmt *S) : S(S) {}
 };
 }
 
@@ -789,13 +835,13 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   auto *AllocBB = createBasicBlock("coro.alloc");
   auto *InitBB = createBasicBlock("coro.init");
   auto *FinalBB = createBasicBlock("coro.final");
-  auto *RetBB = createBasicBlock("coro.ret");
+  auto *EndBB = createBasicBlock("coro.end");
 
   auto *CoroId = Builder.CreateCall(
       CGM.getIntrinsic(llvm::Intrinsic::coro_id),
       {Builder.getInt32(NewAlign), NullPtr, NullPtr, NullPtr});
   createCoroData(*this, CurCoro, CoroId);
-  CurCoro.Data->SuspendBB = RetBB;
+  CurCoro.Data->SuspendBB = EndBB;
   assert(ShouldEmitLifetimeMarkers &&
          "Must emit lifetime intrinsics for coroutines");
 
@@ -831,23 +877,25 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
 
   EmitBlock(InitBB);
 
-  // Pass the result of the allocation to coro.begin.
-  auto *Phi = Builder.CreatePHI(VoidPtrTy, 2);
-  Phi->addIncoming(NullPtr, EntryBB);
-  Phi->addIncoming(AllocateCall, AllocOrInvokeContBB);
-  auto *CoroBegin = Builder.CreateCall(
-      CGM.getIntrinsic(llvm::Intrinsic::coro_begin), {CoroId, Phi});
-  CurCoro.Data->CoroBegin = CoroBegin;
+  {
+    // Pass the result of the allocation to coro.begin.
+    auto *Phi = Builder.CreatePHI(VoidPtrTy, 2);
+    Phi->addIncoming(NullPtr, EntryBB);
+    Phi->addIncoming(AllocateCall, AllocOrInvokeContBB);
+    auto *CoroBegin = Builder.CreateCall(
+        CGM.getIntrinsic(llvm::Intrinsic::coro_begin), {CoroId, Phi});
+    CurCoro.Data->CoroBegin = CoroBegin;
+  }
 
   GetReturnObjectManager GroManager(*this, S);
   GroManager.EmitGroAlloca();
 
-  CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
+  CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(EndBB);
   {
     CGDebugInfo *DI = getDebugInfo();
     ParamReferenceReplacerRAII ParamReplacer(LocalDeclMap);
     CodeGenFunction::RunCleanupsScope ResumeScope(*this);
-    EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, S.getDeallocate());
+    EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, &S);
 
     // Create mapping between parameters and copy-params for coroutine function.
     llvm::ArrayRef<const Stmt *> ParamMoves = S.getParamMoves();
@@ -952,7 +1000,14 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     }
   }
 
-  EmitBlock(RetBB);
+  EmitBlock(EndBB);
+  auto *Phi = Builder.CreatePHI(Builder.getInt1Ty(), llvm::pred_size(EndBB),
+                                "never.suspend");
+  BasicBlock *CleanupBB = CurCoro.Data->InRamp->getParent();
+  for (auto *Pred : llvm::predecessors(EndBB)) {
+    auto *V = (Pred == CleanupBB) ? Builder.getTrue() : Builder.getFalse();
+    Phi->addIncoming(V, Pred);
+  }
   // Emit coro.end before getReturnStmt (and parameter destructors), since
   // resume and destroy parts of the coroutine should not include them.
   llvm::Function *CoroEnd = CGM.getIntrinsic(llvm::Intrinsic::coro_end);
@@ -973,7 +1028,21 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     // shouldn't change the AST.
     if (PreviousRetValue)
       cast<ReturnStmt>(Ret)->setRetValue(PreviousRetValue);
-  }
+
+    // Emit cleanup for ramp function
+    auto *RampCleanupBB = CurCoro.Data->RampCleanupBB;
+    auto *RetBB = EndBB->splitBasicBlock(EndBB->getTerminator(), "coro.ret");
+    EndBB->getTerminator()->eraseFromParent();
+
+    Builder.SetInsertPoint(EndBB);
+    Builder.CreateCondBr(Phi, RampCleanupBB, RetBB);
+
+    EmitBlock(RampCleanupBB);
+    CallCoroDelete(&S).EmitCoroFree(*this, ".ramp");
+    Builder.CreateBr(RetBB);
+    Builder.ClearInsertionPoint();
+  } else
+    EndBB->setName("coro.ret");
 
   // LLVM require the frontend to mark the coroutine.
   CurFn->setPresplitCoroutine();
diff --git a/clang/test/CodeGenCXX/type-aware-coroutines.cpp b/clang/test/CodeGenCXX/type-aware-coroutines.cpp
@@ -55,7 +55,7 @@ extern "C" resumable f1(int) {
   co_return;
 // CHECK: coro.alloc:
 // CHECK: _ZN9resumable12promise_typenwEmi
-// CHECK: coro.free:
+// CHECK: coro.free.ramp:
 // CHECK: _ZN9resumable12promise_typedlEPv
 }
 
@@ -64,7 +64,7 @@ extern "C" resumable f2(float) {
   co_return;
 // CHECK: coro.alloc:
 // CHECK: _ZN9resumable12promise_typenwEmf
-// CHECK: coro.free:
+// CHECK: coro.free.ramp:
 // CHECK: _ZN9resumable12promise_typedlEPv
 }
 
@@ -74,7 +74,7 @@ extern "C" resumable2 f3(int, float, const char*, Allocator) {
    co_return;
 // CHECK: coro.alloc:
 // CHECK: _ZN10resumable212promise_typenwIJifPKc9AllocatorEEEPvmDpT_
-// CHECK: coro.free:
+// CHECK: coro.free.ramp:
 // CHECK: _ZN10resumable212promise_typedlEPv
 // CHECK: _ZN10resumable212promise_typedlEPv
 }
@@ -84,7 +84,7 @@ extern "C" resumable f4(int n = 10) {
    for (int i = 0; i < n; i++) co_yield i;
 // CHECK: coro.alloc:
 // CHECK: call {{.*}}@_ZN9resumable12promise_typenwEmi(
-// CHECK: coro.free:
+// CHECK: coro.free.ramp:
 // CHECK: call void @_ZN9resumable12promise_typedlEPv(
 // CHECK: call void @_ZN9resumable12promise_typedlEPv(
 }
@@ -110,7 +110,7 @@ extern "C" resumable3 f5(float) {
   co_return;
 // CHECK: coro.alloc:
 // CHECK: call {{.*}}@_Znwm(
-// CHECK: coro.free:
+// CHECK: coro.free.ramp:
 // CHECK: call void @_ZdlPvm(
 // CHECK: call void @_ZdlPvm(
 }
diff --git a/clang/test/CodeGenCoroutines/coro-dest-slot.cpp b/clang/test/CodeGenCoroutines/coro-dest-slot.cpp
@@ -25,7 +25,7 @@ extern "C" coro f(int) { co_return; }
 // CHECK: %[[CLEANUP_DEST0:.+]] = phi i32 [ 0, %[[INIT_READY]] ], [ 2, %[[INIT_CLEANUP]] ]
 
 // CHECK: %[[FINAL_SUSPEND:.+]] = call i8 @llvm.coro.suspend(
-// CHECK-NEXT: switch i8 %{{.*}}, label %coro.ret [
+// CHECK-NEXT: switch i8 %{{.*}}, label %coro.end [
 // CHECK-NEXT:   i8 0, label %[[FINAL_READY:.+]]
 // CHECK-NEXT:   i8 1, label %[[FINAL_CLEANUP:.+]]
 // CHECK-NEXT: ]
diff --git a/clang/test/CodeGenCoroutines/coro-promise-dtor2.cpp b/clang/test/CodeGenCoroutines/coro-promise-dtor2.cpp
@@ -0,0 +1,57 @@
+// Test that we do not destroy the coroutine promise before coroutine completion, even if the coroutine never suspends.
+// RUN: %clang_cc1 -std=c++20 -triple=x86_64-unknown-linux-gnu -emit-llvm -o - %s \
+// RUN:            -fexceptions -fcxx-exceptions -disable-llvm-passes | FileCheck %s --check-prefix=CHECK-O0
+// RUN: %clang_cc1 -std=c++20 -triple=x86_64-unknown-linux-gnu -emit-llvm -o - %s \
+// RUN:            -fexceptions -fcxx-exceptions -O2 | FileCheck %s --check-prefix=CHECK-O2
+
+#include "Inputs/coroutine.h"
+
+template<class T>
+struct coro {
+  struct RValueWrapper {
+    T* p;
+
+    operator T&&() const noexcept { return static_cast<T&&>(*p); }
+  };
+
+  using promise_type = T;
+
+  T& getDerived() noexcept { return *static_cast<T*>(this); }
+
+  auto get_return_object() noexcept { return RValueWrapper(&getDerived()); }
+  std::suspend_never initial_suspend() noexcept { return {}; }
+  std::suspend_never final_suspend() noexcept { return {}; }
+  void return_value(T&& x) noexcept { getDerived() = static_cast<T&&>(x); }
+  void unhandled_exception() {}
+};
+
+struct A : public coro<A> {
+  int a;
+
+  ~A() {}
+};
+
+A func() {
+  A aa{};
+  aa.a = 5;
+  co_return static_cast<A&&>(aa);
+}
+
+// CHECK-O0-LABEL: define {{.*}} void @_Z4funcv(
+// CHECK-O0: coro.end:
+// CHECK-O0-NEXT: %never.suspend = phi i1 [ false, %cleanup.cont16 ], [ true, %cleanup12 ], [ false, %after.coro.free ], [ false, %final.suspend ], [ false, %init.suspend ]
+// CHECK-O0-NEXT: call void @llvm.coro.end
+// CHECK-O0-NEXT: ptr @_ZNK4coroI1AE13RValueWrappercvOS0_Ev
+// CHECK-O0-NEXT: call void @llvm.memcpy.p0.p0.i64
+// CHECK-O0-NEXT: br i1 %never.suspend, label %coro.cleanup.ramp, label %coro.ret
+
+// CHECK-O0: coro.cleanup.ramp:
+// CHECK-O0-NEXT: call void @_ZN1AD1Ev
+// CHECK-O0-NEXT: call void @llvm.lifetime.end.p0
+// CHECK-O0-NEXT: call ptr @llvm.coro.free
+
+// CHECK-O0: coro.ret:
+// CHECK-O0-NEXT: call void @llvm.lifetime.end.p0
+// CHECK-O0-NEXT: ret void
+
+// CHECK-O2: store i32 5
diff --git a/clang/test/CodeGenCoroutines/coro-symmetric-transfer-02.cpp b/clang/test/CodeGenCoroutines/coro-symmetric-transfer-02.cpp
@@ -91,7 +91,7 @@ Task bar() {
 // CHECK-NEXT:    %{{.+}} = call token @llvm.coro.save(ptr null)
 // CHECK-NEXT:    call void @llvm.coro.await.suspend.handle
 // CHECK-NEXT:    %{{.+}} = call i8 @llvm.coro.suspend
-// CHECK-NEXT:    switch i8 %{{.+}}, label %coro.ret [
+// CHECK-NEXT:    switch i8 %{{.+}}, label %coro.end [
 // CHECK-NEXT:      i8 0, label %[[CASE1_AWAIT_READY]]
 // CHECK-NEXT:      i8 1, label %[[CASE1_AWAIT_CLEANUP:.+]]
 // CHECK-NEXT:    ]
@@ -105,7 +105,7 @@ Task bar() {
 // CHECK-NEXT:    %{{.+}} = call token @llvm.coro.save(ptr null)
 // CHECK-NEXT:    call void @llvm.coro.await.suspend.handle
 // CHECK-NEXT:    %{{.+}} = call i8 @llvm.coro.suspend
-// CHECK-NEXT:    switch i8 %{{.+}}, label %coro.ret [
+// CHECK-NEXT:    switch i8 %{{.+}}, label %coro.end [
 // CHECK-NEXT:      i8 0, label %[[CASE2_AWAIT_READY]]
 // CHECK-NEXT:      i8 1, label %[[CASE2_AWAIT_CLEANUP:.+]]
 // CHECK-NEXT:    ]
