[clang] ce7eb2e - [Coroutines] Avoid creating conditional cleanup markers in suspend block

[Wei Wang via cfe-commits]

Author: Wei Wang
Date: 2023-02-28T15:30:04-08:00
New Revision: ce7eb2e05544437af127684030a21b9e54a34f93

URL: https://github.com/llvm/llvm-project/commit/ce7eb2e05544437af127684030a21b9e54a34f93
DIFF: https://github.com/llvm/llvm-project/commit/ce7eb2e05544437af127684030a21b9e54a34f93.diff

LOG: [Coroutines] Avoid creating conditional cleanup markers in suspend block

We shouldn't access coro frame after returning from `await_suspend()` and before `llvm.coro.suspend()`.
Make sure we always hoist conditional cleanup markers when inside the `await.suspend` block.

Fix https://github.com/llvm/llvm-project/issues/59181

Reviewed By: ChuanqiXu

Differential Revision: https://reviews.llvm.org/D144680

Added: 
    clang/test/CodeGenCoroutines/pr59181.cpp

Modified: 
    clang/lib/CodeGen/CGCoroutine.cpp
    clang/lib/CodeGen/CGExpr.cpp
    clang/lib/CodeGen/CodeGenFunction.h

Removed: 
    


################################################################################
diff  --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 775a4341558a..9b233c1807cf 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -198,7 +198,9 @@ static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Co
   auto *NullPtr = llvm::ConstantPointerNull::get(CGF.CGM.Int8PtrTy);
   auto *SaveCall = Builder.CreateCall(CoroSave, {NullPtr});
 
+  CGF.CurCoro.InSuspendBlock = true;
   auto *SuspendRet = CGF.EmitScalarExpr(S.getSuspendExpr());
+  CGF.CurCoro.InSuspendBlock = false;
   if (SuspendRet != nullptr && SuspendRet->getType()->isIntegerTy(1)) {
     // Veto suspension if requested by bool returning await_suspend.
     BasicBlock *RealSuspendBlock =

diff  --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp
index 052701c1ea18..a9116dd9c5b0 100644
--- a/clang/lib/CodeGen/CGExpr.cpp
+++ b/clang/lib/CodeGen/CGExpr.cpp
@@ -541,13 +541,17 @@ EmitMaterializeTemporaryExpr(const MaterializeTemporaryExpr *M) {
       // Avoid creating a conditional cleanup just to hold an llvm.lifetime.end
       // marker. Instead, start the lifetime of a conditional temporary earlier
       // so that it's unconditional. Don't do this with sanitizers which need
-      // more precise lifetime marks.
+      // more precise lifetime marks. However when inside an "await.suspend"
+      // block, we should always avoid conditional cleanup because it creates
+      // boolean marker that lives across await_suspend, which can destroy coro
+      // frame.
       ConditionalEvaluation *OldConditional = nullptr;
       CGBuilderTy::InsertPoint OldIP;
       if (isInConditionalBranch() && !E->getType().isDestructedType() &&
-          !SanOpts.has(SanitizerKind::HWAddress) &&
-          !SanOpts.has(SanitizerKind::Memory) &&
-          !CGM.getCodeGenOpts().SanitizeAddressUseAfterScope) {
+          ((!SanOpts.has(SanitizerKind::HWAddress) &&
+            !SanOpts.has(SanitizerKind::Memory) &&
+            !CGM.getCodeGenOpts().SanitizeAddressUseAfterScope) ||
+           inSuspendBlock())) {
         OldConditional = OutermostConditional;
         OutermostConditional = nullptr;
 

diff  --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index d7fef7ce5482..4298eb6c2b71 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -333,6 +333,7 @@ class CodeGenFunction : public CodeGenTypeCache {
   // in this header.
   struct CGCoroInfo {
     std::unique_ptr<CGCoroData> Data;
+    bool InSuspendBlock = false;
     CGCoroInfo();
     ~CGCoroInfo();
   };
@@ -342,6 +343,10 @@ class CodeGenFunction : public CodeGenTypeCache {
     return CurCoro.Data != nullptr;
   }
 
+  bool inSuspendBlock() const {
+    return isCoroutine() && CurCoro.InSuspendBlock;
+  }
+
   /// CurGD - The GlobalDecl for the current function being compiled.
   GlobalDecl CurGD;
 

diff  --git a/clang/test/CodeGenCoroutines/pr59181.cpp b/clang/test/CodeGenCoroutines/pr59181.cpp
new file mode 100644
index 000000000000..80f4634db252
--- /dev/null
+++ b/clang/test/CodeGenCoroutines/pr59181.cpp
@@ -0,0 +1,60 @@
+// Test for PR59181. Tests that no conditional cleanup is created around await_suspend.
+//
+// REQUIRES: x86-registered-target
+//
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -emit-llvm %s -o - -std=c++20 -disable-llvm-passes -fsanitize-address-use-after-scope | FileCheck %s
+
+#include "Inputs/coroutine.h"
+
+struct Task {
+  int value_;
+  struct promise_type {
+    Task get_return_object() {
+      return Task{0};
+    }
+
+    std::suspend_never initial_suspend() noexcept {
+      return {};
+    }
+
+    std::suspend_never final_suspend() noexcept {
+      return {};
+    }
+
+    void return_value(Task t) noexcept {}
+    void unhandled_exception() noexcept {}
+
+    auto await_transform(Task t) {
+      struct Suspension {
+        auto await_ready() noexcept { return false;}
+        auto await_suspend(std::coroutine_handle<> coro) {
+          coro.destroy();
+        }
+
+        auto await_resume() noexcept {
+          return 0;
+        }
+      };
+      return Suspension{};
+    }
+  };
+};
+
+Task bar(bool cond) {
+  co_return cond ? Task{ co_await Task{}}: Task{};
+}
+
+void foo() {
+  bar(true);
+}
+
+// CHECK: cleanup.cont:{{.*}}
+// CHECK-NEXT: load i8
+// CHECK-NEXT: trunc
+// CHECK-NEXT: store i1 false
+// CHECK-NEXT: call void @llvm.lifetime.start.p0(i64 8, ptr [[REF:%ref.tmp[0-9]+]])
+
+// CHECK: await.suspend:{{.*}}
+// CHECK-NOT: call void @llvm.lifetime.start.p0(i64 8, ptr [[REF]])
+// CHECK: call void @_ZZN4Task12promise_type15await_transformES_EN10Suspension13await_suspendESt16coroutine_handleIvE
+// CHECK-NEXT: call void @llvm.lifetime.end.p0(i64 8, ptr [[REF]])