[coroutines] Add co_return statement emission

Summary:
Added co_return statement emission.

Tweaked coro-alloc.cpp test to use co_return to trigger coroutine processing instead of co_await, since this change starts emitting the body of the coroutine and await expression handling has not been upstreamed yet.

Reviewers: rsmith, majnemer, EricWF, aaron.ballman

Reviewed By: rsmith

Subscribers: majnemer, llvm-commits, mehdi_amini

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

llvm-svn: 297076
diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 2fdb127..77497f3 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -21,6 +21,13 @@
 namespace CodeGen {
 
 struct CGCoroData {
+
+  // Stores the jump destination just before the final suspend. Coreturn
+  // statements jumps to this point after calling return_xxx promise member.
+  CodeGenFunction::JumpDest FinalJD;
+
+  unsigned CoreturnCount = 0;
+
   // Stores the llvm.coro.id emitted in the function so that we can supply it
   // as the first argument to coro.begin, coro.alloc and coro.free intrinsics.
   // Note: llvm.coro.id returns a token that cannot be directly expressed in a
@@ -59,19 +66,46 @@
   CurCoro.Data->CoroIdExpr = CoroIdExpr;
 }
 
+void CodeGenFunction::EmitCoreturnStmt(CoreturnStmt const &S) {
+  ++CurCoro.Data->CoreturnCount;
+  EmitStmt(S.getPromiseCall());
+  EmitBranchThroughCleanup(CurCoro.Data->FinalJD);
+}
+
 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 *CoroId = Builder.CreateCall(
       CGM.getIntrinsic(llvm::Intrinsic::coro_id),
       {Builder.getInt32(NewAlign), NullPtr, NullPtr, NullPtr});
   createCoroData(*this, CurCoro, CoroId);
 
   EmitScalarExpr(S.getAllocate());
-  // FIXME: Emit the rest of the coroutine.
+
+  // FIXME: Setup cleanup scopes.
+
+  EmitStmt(S.getPromiseDeclStmt());
+
+  CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
+
+  // FIXME: Emit initial suspend and more before the body.
+
+  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());
+
+  // FIXME: Emit return for the coroutine return object.
 }
 
 // Emit coroutine intrinsic and patch up arguments of the token type.
diff --git a/clang/lib/CodeGen/CGStmt.cpp b/clang/lib/CodeGen/CGStmt.cpp
index 8370607..9c92367 100644
--- a/clang/lib/CodeGen/CGStmt.cpp
+++ b/clang/lib/CodeGen/CGStmt.cpp
@@ -145,7 +145,7 @@
     EmitCoroutineBody(cast<CoroutineBodyStmt>(*S));
     break;
   case Stmt::CoreturnStmtClass:
-    CGM.ErrorUnsupported(S, "coroutine");
+    EmitCoreturnStmt(cast<CoreturnStmt>(*S));
     break;
   case Stmt::CapturedStmtClass: {
     const CapturedStmt *CS = cast<CapturedStmt>(S);
diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h
index 347908b..659ec0f 100644
--- a/clang/lib/CodeGen/CodeGenFunction.h
+++ b/clang/lib/CodeGen/CodeGenFunction.h
@@ -2506,6 +2506,7 @@
   void EmitObjCAutoreleasePoolStmt(const ObjCAutoreleasePoolStmt &S);
 
   void EmitCoroutineBody(const CoroutineBodyStmt &S);
+  void EmitCoreturnStmt(const CoreturnStmt &S);
   RValue EmitCoroutineIntrinsic(const CallExpr *E, unsigned int IID);
 
   void EnterCXXTryStmt(const CXXTryStmt &S, bool IsFnTryBlock = false);