[coroutines] Implement correct GRO lifetime

[GorNishanov]

Summary

Sema creates a declaration for gro variable as:

auto $gro = $promise.get_return_object();

However, gro variable has to outlive coroutine frame and coroutine promise, but,
it can only be initialized after the coroutine promise was created, thus, we
split its emission in two parts: EmitGroAlloca emits an alloca and sets up
the cleanups. Later when the coroutine promise is available we initialize
the gro and set the flag that the cleanup is now active.

[rsmith]

lib/CodeGen/CGCoroutine.cpp

366

It doesn't seem safe to assume that a prior EHCleanupScope would be for the GRO variable, nor that it can only have one such cleanup. How about this: grab a stable EHStack iterator before you emit cleanups, grab another afterwards, and iterate over the cleanups in that range modifying them as appropriate.

Even that seems a bit fragile, though. Would it be feasible to thread the 'active' flag through EmitAutoVarCleanups (perhaps add it to AutoVarEmission)?

[GorNishanov]

Remember old_top before emitting the cleanup and walk set Active on all emitted cleanups.
I tried to add Active flag to Emission, but it ended up being very hairy, so I went with the first option suggested.

[GorNishanov]

@rsmith, better now? :)

[GorNishanov]

merged with top of the trunk

[GorNishanov]

merge to top of the trunk

[GorNishanov]

merge with tot, preparing to land

[GorNishanov]

Landed as: https://reviews.llvm.org/rL303716

[GorNishanov]

Closed by commit rL303716: [coroutines] Implement correct GRO lifetime (authored by GorNishanov)

lib/CodeGen/CGCoroutine.cpp

 //===----- CGCoroutine.cpp - Emit LLVM Code for C++ coroutines ------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This contains code dealing with C++ code generation of coroutines.
 //
 //===----------------------------------------------------------------------===//
 
+#include "CGCleanup.h"
 #include "CodeGenFunction.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "clang/AST/StmtCXX.h"
 
 using namespace clang;
 using namespace CodeGen;
 
 using llvm::Value;
void Emit(CodeGenFunction &CGF, Flags) override {
     // No longer need old terminator.
     InsertPt->eraseFromParent();
     CGF.Builder.SetInsertPoint(AfterFreeBB);
   }
   explicit CallCoroDelete(Stmt *DeallocStmt) : Deallocate(DeallocStmt) {}
 };
 }
 
+namespace {
+struct GetReturnObjectManager {
+  CodeGenFunction &CGF;
+  CGBuilderTy &Builder;
+  const CoroutineBodyStmt &S;
+
+  Address GroActiveFlag;
+  CodeGenFunction::AutoVarEmission GroEmission;
+
+  GetReturnObjectManager(CodeGenFunction &CGF, const CoroutineBodyStmt &S)
+      : CGF(CGF), Builder(CGF.Builder), S(S), GroActiveFlag(Address::invalid()),
+        GroEmission(CodeGenFunction::AutoVarEmission::invalid()) {}
+
+  // 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() {
+    auto *GroDeclStmt = dyn_cast<DeclStmt>(S.getResultDecl());
+    if (!GroDeclStmt) {
+      // If get_return_object returns void, no need to do an alloca.
+      return;
+    }
+
+    auto *GroVarDecl = cast<VarDecl>(GroDeclStmt->getSingleDecl());
+
+    // Set GRO flag that it is not initialized yet
+    GroActiveFlag =
+      CGF.CreateTempAlloca(Builder.getInt1Ty(), CharUnits::One(), "gro.active");
+    Builder.CreateStore(Builder.getFalse(), GroActiveFlag);
+
+    GroEmission = CGF.EmitAutoVarAlloca(*GroVarDecl);
+
+    // Remember the top of EHStack before emitting the cleanup.
+    auto old_top = CGF.EHStack.stable_begin();
+    CGF.EmitAutoVarCleanups(GroEmission);
+    auto top = CGF.EHStack.stable_begin();
+
+    // Make the cleanup conditional on gro.active
+    for (auto b = CGF.EHStack.find(top), e = CGF.EHStack.find(old_top);
+      b != e; b++) {
+      if (auto *Cleanup = dyn_cast<EHCleanupScope>(&*b)) {
+        assert(!Cleanup->hasActiveFlag() && "cleanup already has active flag?");
+        Cleanup->setActiveFlag(GroActiveFlag);
+        Cleanup->setTestFlagInEHCleanup();
+        Cleanup->setTestFlagInNormalCleanup();
+      }
+    }
+  }
+
+  void EmitGroInit() {
+    if (!GroActiveFlag.isValid()) {
+      // No Gro variable was allocated. Simply emit the call to
+      // get_return_object.
+      CGF.EmitStmt(S.getResultDecl());
+      return;
+    }
+
+    CGF.EmitAutoVarInit(GroEmission);
+    Builder.CreateStore(Builder.getTrue(), GroActiveFlag);
+  }
+};
+}
+
 static void emitBodyAndFallthrough(CodeGenFunction &CGF,
                                    const CoroutineBodyStmt &S, Stmt *Body) {
   CGF.EmitStmt(Body);
   const bool CanFallthrough = CGF.Builder.GetInsertBlock();
   if (CanFallthrough)
     if (Stmt *OnFallthrough = S.getFallthroughHandler())
       CGF.EmitStmt(OnFallthrough);
 }
void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   // 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);
   {
     CodeGenFunction::RunCleanupsScope ResumeScope(*this);
     EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, S.getDeallocate());
 
     EmitStmt(S.getPromiseDeclStmt());
-    EmitStmt(S.getResultDecl()); // FIXME: Gro lifetime is wrong.
 
+    // Now we have the promise, initialize the GRO
+    GroManager.EmitGroInit();
     EHStack.pushCleanup<CallCoroEnd>(EHCleanup);
 
     CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 
     // FIXME: Emit param moves.
 
     CurCoro.Data->CurrentAwaitKind = AwaitKind::Init;
     EmitStmt(S.getInitSuspendStmt());

test/CodeGenCoroutines/coro-gro.cpp

+// Verifies lifetime of __gro local variable
+// Verify that coroutine promise and allocated memory are freed up on exception.
+// RUN: %clang_cc1 -std=c++1z -fcoroutines-ts -triple=x86_64-unknown-linux-gnu -emit-llvm -o - %s -disable-llvm-passes | FileCheck %s
+
+namespace std::experimental {
+template <typename... T> struct coroutine_traits;
+
+template <class Promise = void> struct coroutine_handle {
+  coroutine_handle() = default;
+  static coroutine_handle from_address(void *) noexcept;
+};
+template <> struct coroutine_handle<void> {
+  static coroutine_handle from_address(void *) noexcept;
+  coroutine_handle() = default;
+  template <class PromiseType>
+  coroutine_handle(coroutine_handle<PromiseType>) noexcept;
+};
+}
+
+struct suspend_always {
+  bool await_ready() noexcept;
+  void await_suspend(std::experimental::coroutine_handle<>) noexcept;
+  void await_resume() noexcept;
+};
+
+struct GroType {
+  ~GroType();
+  operator int() noexcept;
+};
+
+template <> struct std::experimental::coroutine_traits<int> {
+  struct promise_type {
+    GroType get_return_object() noexcept;
+    suspend_always initial_suspend() noexcept;
+    suspend_always final_suspend() noexcept;
+    void return_void() noexcept;
+    promise_type();
+    ~promise_type();
+    void unhandled_exception() noexcept;
+  };
+};
+
+struct Cleanup { ~Cleanup(); };
+void doSomething() noexcept;
+
+// CHECK: define i32 @_Z1fv(
+int f() {
+  // CHECK: %[[RetVal:.+]] = alloca i32
+  // CHECK: %[[GroActive:.+]] = alloca i1
+
+  // CHECK: %[[Size:.+]] = call i64 @llvm.coro.size.i64()
+  // CHECK: call i8* @_Znwm(i64 %[[Size]])
+  // CHECK: store i1 false, i1* %[[GroActive]]
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJiEE12promise_typeC1Ev(
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJiEE12promise_type17get_return_objectEv(
+  // CHECK: store i1 true, i1* %[[GroActive]]
+
+  Cleanup cleanup;
+  doSomething();
+  co_return;
+
+  // CHECK: call void @_Z11doSomethingv(
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJiEE12promise_type11return_voidEv(
+  // CHECK: call void @_ZN7CleanupD1Ev(
+
+  // Destroy promise and free the memory.
+
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJiEE12promise_typeD1Ev(
+  // CHECK: %[[Mem:.+]] = call i8* @llvm.coro.free(
+  // CHECK: call void @_ZdlPv(i8* %[[Mem]])
+
+  // Initialize retval from Gro and destroy Gro
+
+  // CHECK: %[[Conv:.+]] = call i32 @_ZN7GroTypecviEv(
+  // CHECK: store i32 %[[Conv]], i32* %[[RetVal]]
+  // CHECK: %[[IsActive:.+]] = load i1, i1* %[[GroActive]]
+  // CHECK: br i1 %[[IsActive]], label %[[CleanupGro:.+]], label %[[Done:.+]]
+
+  // CHECK: [[CleanupGro]]:
+  // CHECK:   call void @_ZN7GroTypeD1Ev(
+  // CHECK:   br label %[[Done]]
+
+  // CHECK: [[Done]]:
+  // CHECK:   %[[LoadRet:.+]] = load i32, i32* %[[RetVal]]
+  // CHECK:   ret i32 %[[LoadRet]]
+}