[coroutines] Add cleanup for compiler injected objects/allocations in coroutine body

Summary:
* Use pushCleanup to emit freeing coroutine memory on normal and EH exits.
* Surround emitted code with CodeGenFunction::RunCleanupsScope.

Reviewers: rsmith, rnk, EricWF

Reviewed By: rnk

Subscribers: cfe-commits

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

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@299281 91177308-0d34-0410-b5e6-96231b3b80d8

Author: GorNishanov

lib/CodeGen/CGCoroutine.cpp

@@ -215,6 +215,24 @@ void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
   EmitBranchThroughCleanup(CurCoro.Data->FinalJD);
 }
 
+namespace {
+// Make sure to call coro.delete on scope exit.
+struct CallCoroDelete final : public EHScopeStack::Cleanup {
+  Stmt *Deallocate;
+
+  // TODO: Wrap deallocate in if(coro.free(...)) Deallocate.
+  void Emit(CodeGenFunction &CGF, Flags) override {
+    // Note: That deallocation will be emitted twice: once for a normal exit and
+    // once for exceptional exit. This usage is safe because Deallocate does not
+    // contain any declarations. The SubStmtBuilder::makeNewAndDeleteExpr()
+    // builds a single call to a deallocation function which is safe to emit
+    // multiple times.
+    CGF.EmitStmt(Deallocate);
+  }
+  explicit CallCoroDelete(Stmt *DeallocStmt) : Deallocate(DeallocStmt) {}
+};
+}
+
 void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
   auto *NullPtr = llvm::ConstantPointerNull::get(Builder.getInt8PtrTy());
   auto &TI = CGM.getContext().getTargetInfo();
@@ -248,26 +266,28 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     EmitBlock(InitBB);
   }
 
-  // FIXME: Setup cleanup scopes.
+  CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
+  {
+    CodeGenFunction::RunCleanupsScope ResumeScope(*this);
+    EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, S.getDeallocate());
 
-  EmitStmt(S.getPromiseDeclStmt());
+    EmitStmt(S.getPromiseDeclStmt());
 
-  CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
-  CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
+    CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 
-  // FIXME: Emit initial suspend and more before the body.
+    // FIXME: Emit initial suspend and more before the body.
 
-  CurCoro.Data->CurrentAwaitKind = AwaitKind::Normal;
-  EmitStmt(S.getBody());
+    CurCoro.Data->CurrentAwaitKind = AwaitKind::Normal;
+    EmitStmt(S.getBody());
 
-  // See if we need to generate final suspend.
-  const bool CanFallthrough = Builder.GetInsertBlock();
-  const bool HasCoreturns = CurCoro.Data->CoreturnCount > 0;
-  if (CanFallthrough || HasCoreturns) {
-    EmitBlock(FinalBB);
-    // FIXME: Emit final suspend.
+    // See if we need to generate final suspend.
+    const bool CanFallthrough = Builder.GetInsertBlock();
+    const bool HasCoreturns = CurCoro.Data->CoreturnCount > 0;
+    if (CanFallthrough || HasCoreturns) {
+      EmitBlock(FinalBB);
+      // FIXME: Emit final suspend.
+    }
   }
-  EmitStmt(S.getDeallocate());
 
   EmitBlock(RetBB);
 

test/CodeGenCoroutines/coro-cleanup.cpp

@@ -0,0 +1,74 @@
+// 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 -fexceptions -fcxx-exceptions -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 *) { return {}; }
+};
+template <> struct coroutine_handle<void> {
+  static coroutine_handle from_address(void *) { return {}; }
+  coroutine_handle() = default;
+  template <class PromiseType>
+  coroutine_handle(coroutine_handle<PromiseType>) {}
+};
+}
+
+struct suspend_always {
+  bool await_ready();
+  void await_suspend(std::experimental::coroutine_handle<>);
+  void await_resume();
+};
+
+template <> struct std::experimental::coroutine_traits<void> {
+  struct promise_type {
+    void get_return_object();
+    suspend_always initial_suspend();
+    suspend_always final_suspend();
+    void return_void();
+    promise_type();
+    ~promise_type();
+    void unhandled_exception();
+  };
+};
+
+struct Cleanup { ~Cleanup(); };
+void may_throw();
+
+// CHECK: define void @_Z1fv(
+void f() {
+  // CHECK: call i8* @_Znwm(i64
+
+  // If promise constructor throws, check that we free the memory.
+
+  // CHECK: invoke void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_typeC1Ev(
+  // CHECK-NEXT: to label %{{.+}} unwind label %[[DeallocPad:.+]]
+
+  Cleanup cleanup;
+  may_throw();
+
+  // if may_throw throws, check that we destroy the promise and free the memory.
+
+  // CHECK: invoke void @_Z9may_throwv(
+  // CHECK-NEXT: to label %{{.+}} unwind label %[[PromDtorPad:.+]]
+
+  // CHECK: [[DeallocPad]]:
+  // CHECK-NEXT: landingpad
+  // CHECK-NEXT:   cleanup
+  // CHECK: br label %[[Dealloc:.+]]
+
+  // CHECK: [[PromDtorPad]]:
+  // CHECK-NEXT: landingpad
+  // CHECK-NEXT:   cleanup
+  // CHECK: call void @_ZN7CleanupD1Ev(%struct.Cleanup*
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_typeD1Ev(
+  // CHECK: br label %[[Dealloc]]
+
+  // CHECK: [[Dealloc]]:
+  // CHECK:   %[[Mem:.+]] = call i8* @llvm.coro.free(
+  // CHECK:   call void @_ZdlPv(i8* %[[Mem]])
+
+  co_return;
+}

test/CodeGenCoroutines/coro-return.cpp

@@ -1,25 +1,18 @@
-// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fcoroutines-ts -std=c++14 -emit-llvm %s -o - -disable-llvm-passes | FileCheck %s
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fcoroutines-ts -std=c++1z -emit-llvm %s -o - -disable-llvm-passes | FileCheck %s
 
-namespace std {
-namespace experimental {
-template <typename... T>
-struct coroutine_traits;
+namespace std::experimental {
+template <typename... T> struct coroutine_traits;
 
-template <class Promise = void>
-struct coroutine_handle {
+template <class Promise = void> struct coroutine_handle {
   coroutine_handle() = default;
   static coroutine_handle from_address(void *) { return {}; }
 };
-
-template <>
-struct coroutine_handle<void> {
+template <> struct coroutine_handle<void> {
   static coroutine_handle from_address(void *) { return {}; }
   coroutine_handle() = default;
   template <class PromiseType>
   coroutine_handle(coroutine_handle<PromiseType>) {}
 };
-
-}
 }
 
 struct suspend_always {
@@ -28,8 +21,7 @@ struct suspend_always {
   void await_resume();
 };
 
-template<>
-struct std::experimental::coroutine_traits<void> {
+template <> struct std::experimental::coroutine_traits<void> {
   struct promise_type {
     void get_return_object();
     suspend_always initial_suspend();