[Coroutines] Avoid creating conditional cleanup markers in suspend block

[weiwang]

Summary

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

[bruno]

Thanks for working on this Wei!

Can you please add a testcase?

clang/lib/CodeGen/CodeGenFunction.h

336

Should this live inside CGCoroData instead?

[ChuanqiXu]

The test should be required here.

clang/lib/CodeGen/CGExpr.cpp

544–547

We should add comment to explain why we add a forward path for coroutines.

[ChuanqiXu]

BTW, what is the conclusion for the concern about sanitizers? Would change make sanitizers to perform false positive diagnostics?

[weiwang]

In D144680#4149149, @ChuanqiXu wrote:

BTW, what is the conclusion for the concern about sanitizers? Would change make sanitizers to perform false positive diagnostics?

This change is limited to the await.suspend block, which only does codegen for awaiter.await_suspend(coroutine_handle<p>::from_promise(p)). In this case, I think the only asan check removed by the change is the conditional marker for cleanup the temporary coroutine_handler used as the parameter of await_suspend, so my understanding is that it shouldn't affect asan diagnostic.

To show the difference it makes to the source from issue #59181

Before:

%ref.tmp25 = alloca %"struct.std::__1::coroutine_handle.6", align 8
...
// asan check inserted here
bool cond_cleanup_marker = false;

if (cond) {
    // get awaiter
    ...

    if (!awaiter.await_ready()) {
        call void @llvm.lifetime.start.p0(i64 8, ptr %ref.tmp25)
        
        // asan check inserted here
	cond_cleanup_marker = true;
	
        // store handler to %ref.tmp25
	...
	
        awaiter.await_suspend();

        // asan check inserted here
	if (cond_cleanup_marker)
	    call void @llvm.lifetime.end.p0(i64 8, ptr %ref.tmp25)
        call i8 @llvm.coro.suspend(...)
	...
    }
    ...
} else {
    ... 
}
...
lpad:
    ...
    // asan check inserted here
    if (cond_cleanup_marker)
        call void @llvm.lifetime.end.p0(i64 8, ptr %ref.tmp25)

The issue is only reproduced in -O0, because `cond_cleanup_marker is optimized away in -O1 or up opt level.

After:

%ref.tmp25 = alloca %"struct.std::__1::coroutine_handle.6", align 8
...

call void @llvm.lifetime.start.p0(i64 8, ptr %ref.tmp25)

if (cond) {
    ...
    if (!awaiter.await_ready()) {
        // store handler to %ref.tmp25
	...

	awaiter.await_suspend();
	call void @llvm.lifetime.end.p0(i64 8, ptr %ref.tmp25)
	call i8 @llvm.coro.suspend(...)
	...
    }
    ...
} else {
    ... 
}
...
lpad:
    call void @llvm.lifetime.end.p0(i64 8, ptr %ref.tmp25)

This also matches with the IR without asan.

clang/lib/CodeGen/CGExpr.cpp

544–547

Will do

clang/lib/CodeGen/CodeGenFunction.h

336

CGCoroData is forward declared here, so it does not know what's inside.

[weiwang]

add test case and some comments

[ChuanqiXu]

LGTM, then.

[weiwang]

add target triple

clang/lib/CodeGen/CGCoroutine.cpp

static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
   // Otherwise, emit suspend logic.
   CGF.EmitBlock(SuspendBlock);
 
   auto &Builder = CGF.Builder;
   llvm::Function *CoroSave = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_save);
   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 =
         CGF.createBasicBlock(Prefix + Twine(".suspend.bool"));
     CGF.Builder.CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock);
     CGF.EmitBlock(RealSuspendBlock);
   }
 

clang/lib/CodeGen/CGExpr.cpp

if (auto *Var = dyn_cast<llvm::GlobalVariable>(
 
     case SD_FullExpression: {
       if (!ShouldEmitLifetimeMarkers)
         break;
 
       // 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::HWAddress) &&
             !SanOpts.has(SanitizerKind::Memory) &&
-          !CGM.getCodeGenOpts().SanitizeAddressUseAfterScope) {
+            !CGM.getCodeGenOpts().SanitizeAddressUseAfterScope) ||
+           inSuspendBlock())) {
         OldConditional = OutermostConditional;
         OutermostConditional = nullptr;
 
         OldIP = Builder.saveIP();
         llvm::BasicBlock *Block = OldConditional->getStartingBlock();
         Builder.restoreIP(CGBuilderTy::InsertPoint(
             Block, llvm::BasicBlock::iterator(Block->back())));
       }

clang/lib/CodeGen/CodeGenFunction.h

public:
   /// Save Parameter Decl for coroutine.
   llvm::SmallVector<const ParmVarDecl *, 4> FnArgs;
 
   // Holds coroutine data if the current function is a coroutine. We use a
   // wrapper to manage its lifetime, so that we don't have to define CGCoroData
   // in this header.
   struct CGCoroInfo {
     std::unique_ptr<CGCoroData> Data;
+    bool InSuspendBlock = false;
     CGCoroInfo();
     ~CGCoroInfo();
   };
   CGCoroInfo CurCoro;
 
   bool isCoroutine() const {
     return CurCoro.Data != nullptr;
   }
 
+  bool inSuspendBlock() const {
+    return isCoroutine() && CurCoro.InSuspendBlock;
+  }
+
   /// CurGD - The GlobalDecl for the current function being compiled.
   GlobalDecl CurGD;
 
   /// PrologueCleanupDepth - The cleanup depth enclosing all the
   /// cleanups associated with the parameters.
   EHScopeStack::stable_iterator PrologueCleanupDepth;
 
   /// ReturnBlock - Unified return block.

clang/test/CodeGenCoroutines/pr59181.cpp

+// 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]])