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

[GorNishanov]

Summary

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

[rnk]

lgtm

lib/CodeGen/CGCoroutine.cpp

225

This will be called twice: once for a normal exit and once for exceptional exit. In general, double emitting a Stmt* is not safe, since it might contain a VarDecl or a LabelDecl, but this usage is safe because SubStmtBuilder::makeNewAndDeleteExpr() builds two calls that can't declare anything. That is *definitely* worth a comment. :)

[GorNishanov]

lib/CodeGen/CGCoroutine.cpp

225

Thank you very much for the review!
Would switching the type of Deallocate from Stmt* to Expr* address this concern?

[rnk]

lib/CodeGen/CGCoroutine.cpp

225

Not really, because gnu statement exprs mean Exprs can contain Decls as well. =(

[GorNishanov]

lib/CodeGen/CGCoroutine.cpp

225

Ahh. Okay, comment it is.

lib/CodeGen/CGCoroutine.cpp

 }
 
 void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
   ++CurCoro.Data->CoreturnCount;
   EmitStmt(S.getPromiseCall());
   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 {
+    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();
   unsigned NewAlign = TI.getNewAlign() / TI.getCharWidth();
 
   auto *FinalBB = createBasicBlock("coro.final");
   auto *RetBB = createBasicBlock("coro.ret");
 
if (auto *RetOnAllocFailure = S.getReturnStmtOnAllocFailure()) {
 
     // If not, return OnAllocFailure object.
     EmitBlock(RetOnFailureBB);
     EmitStmt(RetOnAllocFailure);
 
     EmitBlock(InitBB);
   }
 
-  // FIXME: Setup cleanup scopes.
+  CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
+  {
+    CodeGenFunction::RunCleanupsScope ResumeScope(*this);
+    EHStack.pushCleanup<CallCoroDelete>(NormalAndEHCleanup, S.getDeallocate());
 
     EmitStmt(S.getPromiseDeclStmt());
 
-  CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
     CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 
     // FIXME: Emit initial suspend and more before the body.
 
     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.
     }
-  EmitStmt(S.getDeallocate());
+  }
 
   EmitBlock(RetBB);
 
   // FIXME: Emit return for the coroutine return object.
 }
 
 // Emit coroutine intrinsic and patch up arguments of the token type.
 RValue CodeGenFunction::EmitCoroutineIntrinsic(const CallExpr *E,

test/CodeGenCoroutines/coro-cleanup.cpp

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

-// 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 {
   bool await_ready();
   void await_suspend(std::experimental::coroutine_handle<>);
   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();
     suspend_always final_suspend();
     void return_void();
   };
 };