[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.

Duplicate of: https://reviews.llvm.org/D31670 (which arc patch refuses to apply for some reason)

cfe/trunk/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());

cfe/trunk/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]]
+}