[Clang][Coroutines] Only change cleanup order when GRO does emits an alloca

Author: bcardosolopes

clang/lib/CodeGen/CGCoroutine.cpp

@@ -512,21 +512,28 @@ struct GetReturnObjectManager {
     }();
   }
 
+  bool shouldEmitAlloca() {
+    if (DirectEmit)
+      return false;
+
+    if (!isa_and_nonnull<DeclStmt>(S.getResultDecl())) {
+      // If get_return_object returns void, no need to do an alloca.
+      return false;
+    }
+
+    return true;
+  }
+
   // The gro variable has to outlive coroutine frame and coroutine promise, but,
   // it can only be initialized after coroutine promise was created, thus, we
   // split its emission in two parts. EmitGroAlloca emits an alloca and sets up
   // cleanups. Later when coroutine promise is available we initialize the gro
   // and sets the flag that the cleanup is now active.
   void EmitGroAlloca() {
-    if (DirectEmit)
-      return;
-
-    auto *GroDeclStmt = dyn_cast_or_null<DeclStmt>(S.getResultDecl());
-    if (!GroDeclStmt) {
-      // If get_return_object returns void, no need to do an alloca.
+    if (!shouldEmitAlloca())
       return;
-    }
 
+    auto *GroDeclStmt = cast<DeclStmt>(S.getResultDecl());
     auto *GroVarDecl = cast<VarDecl>(GroDeclStmt->getSingleDecl());
 
     // Set GRO flag that it is not initialized yet
@@ -658,6 +665,25 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
 
   GetReturnObjectManager GroManager(*this, S);
   CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
+  bool shouldGROEmitAlloca = GroManager.shouldEmitAlloca();
+  auto emitRetBlock = [&]() {
+    EmitBlock(RetBB);
+    // 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);
+    Builder.CreateCall(CoroEnd,
+                       {NullPtr, Builder.getFalse(),
+                        llvm::ConstantTokenNone::get(CoroEnd->getContext())});
+
+    if (Stmt *Ret = S.getReturnStmt()) {
+      // Since we already emitted the return value above, so we shouldn't
+      // emit it again here.
+      if (GroManager.DirectEmit)
+        cast<ReturnStmt>(Ret)->setRetValue(nullptr);
+      EmitStmt(Ret);
+    }
+  };
+
   {
     CGDebugInfo *DI = getDebugInfo();
     ParamReferenceReplacerRAII ParamReplacer(LocalDeclMap);
@@ -757,23 +783,16 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
       EmitBlock(FinalBB, /*IsFinished=*/true);
     }
 
-    EmitBlock(RetBB);
-    // 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);
-    Builder.CreateCall(CoroEnd,
-                       {NullPtr, Builder.getFalse(),
-                        llvm::ConstantTokenNone::get(CoroEnd->getContext())});
-
-    if (Stmt *Ret = S.getReturnStmt()) {
-      // Since we already emitted the return value above, so we shouldn't
-      // emit it again here.
-      if (GroManager.DirectEmit)
-        cast<ReturnStmt>(Ret)->setRetValue(nullptr);
-      EmitStmt(Ret);
-    }
+    // If an alloca is emitted for GRO handling, make sure the ret block
+    // is taken into account within the appropriate cleanup.
+    if (shouldGROEmitAlloca)
+      emitRetBlock();
   }
 
+  // No GRO alloca's, no need for to account for extra cleanups.
+  if (!shouldGROEmitAlloca)
+    emitRetBlock();
+
   // LLVM require the frontend to mark the coroutine.
   CurFn->setPresplitCoroutine();
 }

clang/test/CodeGenCoroutines/coro-dest-slot.cpp

@@ -33,22 +33,9 @@ extern "C" coro f(int) { co_return; }
 
 // CHECK: call void @_ZNSt13suspend_never12await_resumeEv(
 // CHECK: %[[CLEANUP_DEST1:.+]] = phi i32 [ 0, %[[FINAL_READY]] ], [ 2, %[[FINAL_CLEANUP]] ]
-// CHECK: switch i32 %[[CLEANUP_DEST1]], label %[[CLEANUP_BB_2:.+]] [
-// CHECK:   i32 0, label %[[CLEANUP_CONT:.+]]
-// CHECK: ]
-
-// CHECK: [[CLEANUP_CONT]]:
-// CHECK:   br label %[[CORO_RET:.+]]
-
-// CHECK: [[CORO_RET]]:
-// CHECK:   call i1 @llvm.coro.end
-// CHECK:   br label %cleanup19
-
-// CHECK: [[CLEANUP_BB_2]]:
-// CHECK: %[[CLEANUP_DEST2:.+]] = phi i32 [ %[[CLEANUP_DEST0]], %{{.+}} ], [ 1, %[[CORO_RET]] ], [ %[[CLEANUP_DEST1]], %{{.+}} ]
-// CHECK: call void @llvm.lifetime.end.p0(i64 1, ptr %[[PROMISE]])
+// CHECK: %[[CLEANUP_DEST2:.+]] = phi i32 [ %[[CLEANUP_DEST0]], %{{.+}} ], [ %[[CLEANUP_DEST1]], %{{.+}} ], [ 0, %{{.+}} ]
 // CHECK: call ptr @llvm.coro.free(
 // CHECK: switch i32 %[[CLEANUP_DEST2]], label %{{.+}} [
+// CHECK-NEXT: i32 0
 // CHECK-NEXT: i32 2
-// CHECK-NEXT: i32 1
 // CHECK-NEXT: ]