[Coroutines] Catch exceptions in await_resume

[modocache]

Summary

http://wg21.link/P0664r2 section "Evolution/Core Issues 24" describes a
proposed change to Coroutines TS that would have any exceptions thrown
after the initial suspend point of a coroutine be caught by the handler
specified by the promise type's 'unhandled_exception' member function.
This commit provides a sample implementation of the specified behavior.

Test Plan: check-clang

[GorNishanov]

Thank you for doing this. It looks very elegant, but, it is a little bit wrong. It creates two different initial_suspend objects.
Run this example:

https://wandbox.org/permlink/Q1Zd2NUlolmw9YmX

You will observe that in trunk we are getting the output:

0x216808c: constructed(initial)
0x216808c: await_ready
0x216808c: await_suspend
0x216808c: await_resume
0x216808c: destroyed
0x21680b8: constructed(final)
0x21680b8: await_ready
0x21680b8: await_suspend
consumed 10 values with sum 45
0x21680b8: destroyed
promise destroyed

With this change, the output becomes:

0x1965c4c: constructed(initial)
0x1965c4c: await_ready
0x1965c4c: await_suspend
0x1965c4c: destroyed
0x1965c60: constructed(initial)
0x1965c60: await_resume
0x1965c60: destroyed
0x1965c80: constructed(final)
0x1965c80: await_ready
0x1965c80: await_suspend
consumed 10 values with sum 45
0x1965c80: destroyed
0x1965c60: destroyed
promise destroyed

I suggest, first modify the unit test to check that we do not create two initial_suspend objects. Then to fix it :-)

lib/CodeGen/CGCoroutine.cpp

620

auto *InitialSuspend = ...

See: https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

636

auto *AwaitExpr =

639

Consider:

std::array<Stmt *,2> Stmts = {AwaitExpr->getResumeExpr(), S.getBody()};

instead

[modocache]

Thanks for the review, @GorNishanov. Here's a more correct solution: an i1 is used to keep track of whether await_resume threw. If it did, the coroutine body is skipped, and we go straight to the final suspend point. Otherwise, the coroutine body is executed as normal.

[GorNishanov]

Looks good. Make sure to run the tests in release and debug mode.
On my box, these three test failed in release mode.

Clang :: CodeGenCoroutines/coro-await-resume-eh.cpp
Clang :: CodeGenCoroutines/coro-promise-dtor.cpp
Clang :: CodeGenCoroutines/coro-unhandled-exception.cpp

There is a small difference in friendly names emission in release mode

lib/CodeGen/CGCoroutine.cpp

220

I suggest to check whether await_resume expression is noexcept and omit emission of all of the goo. Most of the time the await_ready for initial_suspend will be noexcept and thus we would not need to emit extra stuff.

[modocache]

Oops, thanks for testing on release mode, @GorNishanov. Turns out I had a dangling pointer. With this update the tests pass on both release and debug.

[GorNishanov]

LGTM

[modocache]

Thanks again for all the reviews, @GorNishanov! Very much appreciated.

lib/CodeGen/CGCoroutine.cpp

220

Oh shoot, I forgot to do this! I'll take care of this in a follow-up commit within the next week.

lib/CodeGen/CGCoroutine.cpp

struct clang::CodeGen::CGCoroData {
 
   // How many co_return statements are in the coroutine. Used to decide whether
   // we need to add co_return; equivalent at the end of the user authored body.
   unsigned CoreturnCount = 0;
 
   // A branch to this block is emitted when coroutine needs to suspend.
   llvm::BasicBlock *SuspendBB = nullptr;
 
+  // The promise type's 'unhandled_exception' handler, if it defines one.
+  Stmt *ExceptionHandler = nullptr;
+
+  // A temporary i1 alloca that stores whether 'await_resume' threw an
+  // exception. If it did, 'true' is stored in this variable, and the coroutine
+  // body must be skipped. If the promise type does not define an exception
+  // handler, this is null.
+  llvm::Value *ResumeEHVar = nullptr;
+
   // Stores the jump destination just before the coroutine memory is freed.
   // This is the destination that every suspend point jumps to for the cleanup
   // branch.
   CodeGenFunction::JumpDest CleanupJD;
 
   // Stores the jump destination just before the final suspend. The co_return
   // statements jumps to this point after calling return_xxx promise member.
   CodeGenFunction::JumpDest FinalJD;
static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
   Switch->addCase(Builder.getInt8(1), CleanupBlock);
 
   // Emit cleanup for this suspend point.
   CGF.EmitBlock(CleanupBlock);
   CGF.EmitBranchThroughCleanup(Coro.CleanupJD);
 
   // Emit await_resume expression.
   CGF.EmitBlock(ReadyBlock);
+  CXXTryStmt *TryStmt = nullptr;
+  if (Coro.ExceptionHandler && Kind == AwaitKind::Init) {
+    Coro.ResumeEHVar =
+        CGF.CreateTempAlloca(Builder.getInt1Ty(), Prefix + Twine("resume.eh"));
+    Builder.CreateFlagStore(true, Coro.ResumeEHVar);
+
+    auto Loc = S.getResumeExpr()->getExprLoc();
+    auto *Catch = new (CGF.getContext())
+        CXXCatchStmt(Loc, /*exDecl=*/nullptr, Coro.ExceptionHandler);
+    auto *TryBody =
+        CompoundStmt::Create(CGF.getContext(), S.getResumeExpr(), Loc, Loc);
+    TryStmt = CXXTryStmt::Create(CGF.getContext(), Loc, TryBody, Catch);
+    CGF.EnterCXXTryStmt(*TryStmt);
+  }
+
   LValueOrRValue Res;
   if (forLValue)
     Res.LV = CGF.EmitLValue(S.getResumeExpr());
   else
     Res.RV = CGF.EmitAnyExpr(S.getResumeExpr(), aggSlot, ignoreResult);
+
+  if (TryStmt) {
+    Builder.CreateFlagStore(false, Coro.ResumeEHVar);
+    CGF.ExitCXXTryStmt(*TryStmt);
+  }
+
   return Res;
 }
 
 RValue CodeGenFunction::EmitCoawaitExpr(const CoawaitExpr &E,
                                         AggValueSlot aggSlot,
                                         bool ignoreResult) {
   return emitSuspendExpression(*this, *CurCoro.Data, E,
                                CurCoro.Data->CurrentAwaitKind, aggSlot,
CurCoro.Data->CleanupJD = getJumpDestInCurrentScope(RetBB);
     // promise local variable was not emitted yet.
     CoroId->setArgOperand(1, PromiseAddrVoidPtr);
 
     // Now we have the promise, initialize the GRO
     GroManager.EmitGroInit();
 
     EHStack.pushCleanup<CallCoroEnd>(EHCleanup);
 
     CurCoro.Data->CurrentAwaitKind = AwaitKind::Init;
+    CurCoro.Data->ExceptionHandler = S.getExceptionHandler();
     EmitStmt(S.getInitSuspendStmt());
     CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 
     CurCoro.Data->CurrentAwaitKind = AwaitKind::Normal;
 
-    if (auto *OnException = S.getExceptionHandler()) {
+    if (CurCoro.Data->ExceptionHandler) {
+      BasicBlock *BodyBB = createBasicBlock("coro.resumed.body");
+      BasicBlock *ContBB = createBasicBlock("coro.resumed.cont");
+      Value *SkipBody =
+          Builder.CreateFlagLoad(CurCoro.Data->ResumeEHVar, "coro.resumed.eh");
+      Builder.CreateCondBr(SkipBody, ContBB, BodyBB);
+      EmitBlock(BodyBB);
+
       auto Loc = S.getLocStart();
-      CXXCatchStmt Catch(Loc, /*exDecl=*/nullptr, OnException);
-      auto *TryStmt = CXXTryStmt::Create(getContext(), Loc, S.getBody(), &Catch);
+      CXXCatchStmt Catch(Loc, /*exDecl=*/nullptr,
+                         CurCoro.Data->ExceptionHandler);
+      auto *TryStmt =
+          CXXTryStmt::Create(getContext(), Loc, S.getBody(), &Catch);
 
       EnterCXXTryStmt(*TryStmt);
       emitBodyAndFallthrough(*this, S, TryStmt->getTryBlock());
       ExitCXXTryStmt(*TryStmt);
+
+      EmitBlock(ContBB);
     }
     else {
       emitBodyAndFallthrough(*this, S, S.getBody());
     }
 
     // See if we need to generate final suspend.
     const bool CanFallthrough = Builder.GetInsertBlock();
     const bool HasCoreturns = CurCoro.Data->CoreturnCount > 0;

test/CodeGenCoroutines/coro-await-resume-eh.cpp

+// Test the behavior of http://wg21.link/P0664, a proposal to catch any
+// exceptions thrown after the initial suspend point of a coroutine by
+// executing the handler specified by the promise type's 'unhandled_exception'
+// member function.
+//
+// RUN: %clang_cc1 -std=c++14 -fcoroutines-ts \
+// RUN:   -triple=x86_64-unknown-linux-gnu -emit-llvm -o - %s \
+// RUN:   -fexceptions -fcxx-exceptions -disable-llvm-passes \
+// RUN:   | FileCheck %s
+
+#include "Inputs/coroutine.h"
+
+namespace coro = std::experimental::coroutines_v1;
+
+struct throwing_awaitable {
+  bool await_ready() { return true; }
+  void await_suspend(coro::coroutine_handle<>) {}
+  void await_resume() { throw 42; }
+};
+
+struct task {
+  struct promise_type {
+    task get_return_object() { return task{}; }
+    auto initial_suspend() { return throwing_awaitable{}; }
+    auto final_suspend() { return coro::suspend_never{}; }
+    void return_void() {}
+    void unhandled_exception() {}
+  };
+};
+
+// CHECK-LABEL: define void @_Z1fv()
+task f() {
+  // A variable RESUMETHREW is used to keep track of whether the body
+  // of 'await_resume' threw an exception. Exceptions thrown in
+  // 'await_resume' are unwound to RESUMELPAD.
+  // CHECK: init.ready:
+  // CHECK-NEXT: store i1 true, i1* %[[RESUMETHREW:.+]], align 1
+  // CHECK-NEXT: invoke void @_ZN18throwing_awaitable12await_resumeEv
+  // CHECK-NEXT: to label %[[RESUMECONT:.+]] unwind label %[[RESUMELPAD:.+]]
+
+  // If 'await_resume' does not throw an exception, 'false' is stored in
+  // variable RESUMETHREW.
+  // CHECK: [[RESUMECONT]]:
+  // CHECK-NEXT: store i1 false, i1* %[[RESUMETHREW]]
+  // CHECK-NEXT: br label %[[RESUMETRYCONT:.+]]
+
+  // 'unhandled_exception' is called for the exception thrown in
+  // 'await_resume'. The variable RESUMETHREW is never set to false,
+  // and a jump is made to RESUMETRYCONT.
+  // CHECK: [[RESUMELPAD]]:
+  // CHECK: br label %[[RESUMECATCH:.+]]
+  // CHECK: [[RESUMECATCH]]:
+  // CHECK: invoke void @_ZN4task12promise_type19unhandled_exceptionEv
+  // CHECK-NEXT: to label %[[RESUMEENDCATCH:.+]] unwind label
+  // CHECK: [[RESUMEENDCATCH]]:
+  // CHECK-NEXT: invoke void @__cxa_end_catch()
+  // CHECK-NEXT: to label %[[RESUMEENDCATCHCONT:.+]] unwind label
+  // CHECK: [[RESUMEENDCATCHCONT]]:
+  // CHECK-NEXT: br label %[[RESUMETRYCONT]]
+
+  // The variable RESUMETHREW is loaded and if true, then 'await_resume'
+  // threw an exception and the coroutine body is skipped, and the final
+  // suspend is executed immediately. Otherwise, the coroutine body is
+  // executed, and then the final suspend.
+  // CHECK: [[RESUMETRYCONT]]:
+  // CHECK-NEXT: %[[RESUMETHREWLOAD:.+]] = load i1, i1* %[[RESUMETHREW]]
+  // CHECK-NEXT: br i1 %[[RESUMETHREWLOAD]], label %[[RESUMEDCONT:.+]], label %[[RESUMEDBODY:.+]]
+
+  // CHECK: [[RESUMEDBODY]]:
+  // CHECK: invoke void @_ZN4task12promise_type11return_voidEv
+  // CHECK-NEXT: to label %[[REDUMEDBODYCONT:.+]] unwind label
+  // CHECK: [[REDUMEDBODYCONT]]:
+  // CHECK-NEXT: br label %[[COROFINAL:.+]]
+
+  // CHECK: [[RESUMEDCONT]]:
+  // CHECK-NEXT: br label %[[COROFINAL]]
+
+  // CHECK: [[COROFINAL]]:
+  // CHECK-NEXT: invoke void @_ZN4task12promise_type13final_suspendEv
+  co_return;
+}

test/CodeGenCoroutines/coro-unhandled-exception.cpp

 // CHECK:   %[[CATCHSWTOK:.+]] = catchswitch within none [label %[[CATCH:.+]]] unwind label
 // CHECK: [[CATCH]]:
 // CHECK:   %[[CATCHTOK:.+]] = catchpad within [[CATCHSWTOK:.+]]
 // CHECK:   call void @"?unhandled_exception@promise_type@coro_t@@QEAAXXZ"
 // CHECK:   catchret from %[[CATCHTOK]] to label %[[CATCHRETDEST:.+]]
 // CHECK: [[CATCHRETDEST]]:
 // CHECK-NEXT: br label %[[TRYCONT:.+]]
 // CHECK: [[TRYCONT]]:
+// CHECK-NEXT: br label %[[RESUMECONT:.+]]
+// CHECK: [[RESUMECONT]]:
 // CHECK-NEXT: br label %[[COROFIN:.+]]
 // CHECK: [[COROFIN]]:
 // CHECK-NEXT: invoke void @"?final_suspend@promise_type@coro_t@@QEAA?AUsuspend_never@coroutines_v1@experimental@std@@XZ"(
 
 // CHECK-LPAD: @_Z1fv(
 // CHECK-LPAD:   invoke void @_Z9may_throwv()
 // CHECK-LPAD:       to label %[[CONT:.+]] unwind label %[[CLEANUP:.+]]
 // CHECK-LPAD: [[CLEANUP]]:
 // CHECK-LPAD:   call void @_ZN7CleanupD1Ev(%struct.Cleanup* %x) #2
 // CHECK-LPAD:   br label %[[CATCH:.+]]
 
 // CHECK-LPAD: [[CATCH]]:
 // CHECK-LPAD:    call i8* @__cxa_begin_catch
 // CHECK-LPAD:    call void @_ZN6coro_t12promise_type19unhandled_exceptionEv(%"struct.coro_t::promise_type"* %__promise) #2
 // CHECK-LPAD:    invoke void @__cxa_end_catch()
 // CHECK-LPAD-NEXT:  to label %[[CATCHRETDEST:.+]] unwind label
 // CHECK-LPAD: [[CATCHRETDEST]]:
 // CHECK-LPAD-NEXT: br label %[[TRYCONT:.+]]
 // CHECK-LPAD: [[TRYCONT]]:
+// CHECK-LPAD: br label %[[RESUMECONT:.+]]
+// CHECK-LPAD: [[RESUMECONT]]:
 // CHECK-LPAD-NEXT: br label %[[COROFIN:.+]]
 // CHECK-LPAD: [[COROFIN]]:
 // CHECK-LPAD-NEXT: invoke void @_ZN6coro_t12promise_type13final_suspendEv(