[coroutines] Support coroutine-handle returning await-suspend (i.e symmetric control transfer)

[GorNishanov]

Summary

If await_suspend returns a coroutine_handle, as in the example below:

coroutine_handle<> await_suspend(coroutine_handle<> h) {
  coro.promise().waiter = h;
  return coro;
}

suspensionExpression processing will resume the coroutine pointed at by that handle.
Related LLVM change rL311751 makes resume calls of this kind musttail at any optimization level.

This enables unlimited symmetric control transfer from coroutine to coroutine without blowing up the stack.

[GorNishanov]

LGTM

lib/CodeGen/CGCoroutine.cpp

static LValueOrRValue emitSuspendExpression(CodeGenFunction &CGF, CGCoroData &Coro,
   CGF.EmitBlock(SuspendBlock);
 
   auto &Builder = CGF.Builder;
   llvm::Function *CoroSave = CGF.CGM.getIntrinsic(llvm::Intrinsic::coro_save);
   auto *NullPtr = llvm::ConstantPointerNull::get(CGF.CGM.Int8PtrTy);
   auto *SaveCall = Builder.CreateCall(CoroSave, {NullPtr});
 
   auto *SuspendRet = CGF.EmitScalarExpr(S.getSuspendExpr());
-  if (SuspendRet != nullptr) {
+  if (SuspendRet != nullptr && SuspendRet->getType()->isIntegerTy(1)) {
     // Veto suspension if requested by bool returning await_suspend.
-    assert(SuspendRet->getType()->isIntegerTy(1) &&
-           "Sema should have already checked that it is void or bool");
     BasicBlock *RealSuspendBlock =
         CGF.createBasicBlock(Prefix + Twine(".suspend.bool"));
     CGF.Builder.CreateCondBr(SuspendRet, RealSuspendBlock, ReadyBlock);
     SuspendBlock = RealSuspendBlock;
     CGF.EmitBlock(RealSuspendBlock);
   }
 
   // Emit the suspend point.

lib/Sema/SemaCoroutine.cpp

ExprResult Result = S.BuildMemberReferenceExpr(
       SourceLocation(), nullptr, NameInfo, /*TemplateArgs=*/nullptr,
       /*Scope=*/nullptr);
   if (Result.isInvalid())
     return ExprError();
 
   return S.ActOnCallExpr(nullptr, Result.get(), Loc, Args, Loc, nullptr);
 }
 
+// See if return type is coroutine-handle and if so, invoke builtin coro-resume
+// on its address. This is to enable experimental support for coroutine-handle
+// returning await_suspend that results in a guranteed tail call to the target
+// coroutine.
+static Expr *maybeTailCall(Sema &S, QualType RetType, Expr *E,
+                           SourceLocation Loc) {
+  if (RetType->isReferenceType())
+    return nullptr;
+  Type const *T = RetType.getTypePtr();
+  if (!T->isClassType() && !T->isStructureType())
+    return nullptr;
+
+  // FIXME: Add convertability check to coroutine_handle<>. Possibly via
+  // EvaluateBinaryTypeTrait(BTT_IsConvertible, ...) which is at the moment
+  // a private function in SemaExprCXX.cpp
+
+  ExprResult AddressExpr = buildMemberCall(S, E, Loc, "address", None);
+  if (AddressExpr.isInvalid())
+    return nullptr;
+
+  Expr *JustAddress = AddressExpr.get();
+  // FIXME: Check that the type of AddressExpr is void*
+  return buildBuiltinCall(S, Loc, Builtin::BI__builtin_coro_resume,
+                          JustAddress);
+}
+
 /// Build calls to await_ready, await_suspend, and await_resume for a co_await
 /// expression.
 static ReadySuspendResumeResult buildCoawaitCalls(Sema &S, VarDecl *CoroPromise,
                                                   SourceLocation Loc, Expr *E) {
   OpaqueValueExpr *Operand = new (S.Context)
       OpaqueValueExpr(Loc, E->getType(), VK_LValue, E->getObjectKind(), E);
 
   // Assume invalid until we see otherwise.
if (!AwaitReady->getType()->isDependentType()) {
     Calls.Results[ACT::ACT_Ready] = Conv.get();
   }
   CallExpr *AwaitSuspend = cast<CallExpr>(Calls.Results[ACT::ACT_Suspend]);
   if (!AwaitSuspend->getType()->isDependentType()) {
     // [expr.await]p3 [...]
     //   - await-suspend is the expression e.await_suspend(h), which shall be
     //     a prvalue of type void or bool.
     QualType RetType = AwaitSuspend->getCallReturnType(S.Context);
+    // Experimental support for coroutine_handle returning await_suspend.
+    if (Expr *TailCallSuspend = maybeTailCall(S, RetType, AwaitSuspend, Loc))
+      Calls.Results[ACT::ACT_Suspend] = TailCallSuspend;
+    else {
       // non-class prvalues always have cv-unqualified types
       QualType AdjRetType = RetType.getUnqualifiedType();
       if (RetType->isReferenceType() ||
           (AdjRetType != S.Context.BoolTy && AdjRetType != S.Context.VoidTy)) {
         S.Diag(AwaitSuspend->getCalleeDecl()->getLocation(),
                diag::err_await_suspend_invalid_return_type)
             << RetType;
         S.Diag(Loc, diag::note_coroutine_promise_call_implicitly_required)
             << AwaitSuspend->getDirectCallee();
         Calls.IsInvalid = true;
       }
     }
+  }
 
   return Calls;
 }
 
 static ExprResult buildPromiseCall(Sema &S, VarDecl *Promise,
                                    SourceLocation Loc, StringRef Name,
                                    MultiExprArg Args) {
 

test/CodeGenCoroutines/coro-await.cpp

 // RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fcoroutines-ts -std=c++14 \
 // RUN:   -emit-llvm %s -o - -disable-llvm-passes -Wno-coroutine -Wno-unused | FileCheck %s
 
 namespace std {
 namespace experimental {
 template <typename... T>
 struct coroutine_traits;
 
 template <typename Promise = void> struct coroutine_handle;
 
 template <>
 struct coroutine_handle<void> {
   void *ptr;
   static coroutine_handle from_address(void *);
+  void *address();
 };
 
 template <typename Promise>
 struct coroutine_handle : coroutine_handle<> {
   static coroutine_handle from_address(void *);
 };
 
 }
void AwaitReturnsLValue(double) {
   // CHECK: %[[RES2:.+]] = call dereferenceable({{.*}}) %struct.RefTag* @_ZN24AwaitResumeReturnsLValue12await_resumeEv(%struct.AwaitResumeReturnsLValue* %[[TMP1]])
   // CHECK-NEXT: store %struct.RefTag* %[[RES2]], %struct.RefTag** %[[YVAR]],
 
   RefTag& y = co_await AwaitResumeReturnsLValue{};
   // CHECK: %[[RES3:.+]] = call dereferenceable({{.*}}) %struct.RefTag* @_ZN24AwaitResumeReturnsLValue12await_resumeEv(%struct.AwaitResumeReturnsLValue* %[[TMP2]])
   // CHECK-NEXT: store %struct.RefTag* %[[RES3]], %struct.RefTag** %[[ZVAR]],
   RefTag& z = co_yield 42;
 }
+
+struct TailCallAwait {
+  bool await_ready();
+  std::experimental::coroutine_handle<> await_suspend(std::experimental::coroutine_handle<>);
+  void await_resume();
+};
+
+// CHECK-LABEL: @TestTailcall(
+extern "C" void TestTailcall() {
+  co_await TailCallAwait{};
+
+  // CHECK: %[[RESULT:.+]] = call i8* @_ZN13TailCallAwait13await_suspendENSt12experimental16coroutine_handleIvEE(%struct.TailCallAwait*
+  // CHECK: %[[COERCE:.+]] = getelementptr inbounds %"struct.std::experimental::coroutine_handle", %"struct.std::experimental::coroutine_handle"* %[[TMP:.+]], i32 0, i32 0
+  // CHECK: store i8* %[[RESULT]], i8** %[[COERCE]]
+  // CHECK: %[[ADDR:.+]] = call i8* @_ZNSt12experimental16coroutine_handleIvE7addressEv(%"struct.std::experimental::coroutine_handle"* %[[TMP]])
+  // CHECK: call void @llvm.coro.resume(i8* %[[ADDR]])
+}