[coroutines] Add co_return statement emission

[GorNishanov]

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.

[EricWF]

I'm not at all familiar with Clang's codegen mechanics so I won't be able to review this. Maybe @majnemer would?

lib/CodeGen/CGCoroutine.cpp

 
 using namespace clang;
 using namespace CodeGen;
 
 namespace clang {
 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
   // builtin.
   llvm::CallInst *CoroId = nullptr;
   // If coro.id came from the builtin, remember the expression to give better
   // diagnostic. If CoroIdExpr is nullptr, the coro.id was created by
   // EmitCoroutineBody.
if (CurCoro.Data) {
     return;
   }
 
   CurCoro.Data = std::unique_ptr<CGCoroData>(new CGCoroData);
   CurCoro.Data->CoroId = CoroId;
   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.
 RValue CodeGenFunction::EmitCoroutineIntrinsic(const CallExpr *E,
                                                unsigned int IID) {
   SmallVector<llvm::Value *, 8> Args;
   switch (IID) {
   default:

lib/CodeGen/CGStmt.cpp

#include "clang/AST/StmtNodes.inc"
 
   case Stmt::SwitchStmtClass:   EmitSwitchStmt(cast<SwitchStmt>(*S));     break;
   case Stmt::GCCAsmStmtClass:   // Intentional fall-through.
   case Stmt::MSAsmStmtClass:    EmitAsmStmt(cast<AsmStmt>(*S));           break;
   case Stmt::CoroutineBodyStmtClass:
     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);
     EmitCapturedStmt(*CS, CS->getCapturedRegionKind());
     }
     break;
   case Stmt::ObjCAtTryStmtClass:
     EmitObjCAtTryStmt(cast<ObjCAtTryStmt>(*S));

lib/CodeGen/CodeGenFunction.h

#endif
 
   void EmitObjCForCollectionStmt(const ObjCForCollectionStmt &S);
   void EmitObjCAtTryStmt(const ObjCAtTryStmt &S);
   void EmitObjCAtThrowStmt(const ObjCAtThrowStmt &S);
   void EmitObjCAtSynchronizedStmt(const ObjCAtSynchronizedStmt &S);
   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);
   void ExitCXXTryStmt(const CXXTryStmt &S, bool IsFnTryBlock = false);
 
   void EmitCXXTryStmt(const CXXTryStmt &S);
   void EmitSEHTryStmt(const SEHTryStmt &S);
   void EmitSEHLeaveStmt(const SEHLeaveStmt &S);

test/CodeGenCoroutines/coro-alloc.cpp

 extern "C" void f0(global_new_delete_tag) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
-  co_await suspend_always{};
+  co_return;
 }
 
 struct promise_new_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_new_tag> {
   struct promise_type {
     void *operator new(unsigned long);
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }
     suspend_always final_suspend() { return {}; }
     void return_void() {}
   };
 };
 
 // CHECK-LABEL: f1( 
 extern "C" void f1(promise_new_tag ) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_ZNSt12experimental16coroutine_traitsIJv15promise_new_tagEE12promise_typenwEm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
-  co_await suspend_always{};
+  co_return;
 }
 
 struct promise_delete_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_delete_tag> {
   struct promise_type {
     void operator delete(void*);
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }
     suspend_always final_suspend() { return {}; }
     void return_void() {}
   };
 };
 
 // CHECK-LABEL: f2( 
 extern "C" void f2(promise_delete_tag) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJv18promise_delete_tagEE12promise_typedlEPv(i8* %[[MEM]])
-  co_await suspend_always{};
+  co_return;
 }
 
 struct promise_sized_delete_tag {};
 
 template<>
 struct std::experimental::coroutine_traits<void, promise_sized_delete_tag> {
   struct promise_type {
     void operator delete(void*, unsigned long);
     void get_return_object() {}
     suspend_always initial_suspend() { return {}; }
     suspend_always final_suspend() { return {}; }
     void return_void() {}
   };
 };
 
-// CHECK-LABEL: f3( 
+// CHECK-LABEL: f3(
 extern "C" void f3(promise_sized_delete_tag) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: %[[SIZE2:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJv24promise_sized_delete_tagEE12promise_typedlEPvm(i8* %[[MEM]], i64 %[[SIZE2]])
-  co_await suspend_always{};
+  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
+
+namespace std {
+namespace experimental {
+template <typename... T>
+struct coroutine_traits;
+}
+}
+
+struct suspend_always {
+  bool await_ready();
+  void await_suspend();
+  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();
+  };
+};
+
+// CHECK-LABEL: f0(
+extern "C" void f0() {
+  // CHECK: %__promise = alloca %"struct.std::experimental::coroutine_traits<void>::promise_type"
+  // CHECK: %call = call i8* @_Znwm(
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_type11return_voidEv(%"struct.std::experimental::coroutine_traits<void>::promise_type"* %__promise)
+  // CHECK: call void @_ZdlPv
+  co_return;
+}
+
+template<>
+struct std::experimental::coroutine_traits<int> {
+  struct promise_type {
+    int get_return_object();
+    suspend_always initial_suspend();
+    suspend_always final_suspend();
+    void return_value(int);
+  };
+};
+
+// CHECK-LABEL: f1(
+extern "C" int f1() {
+  // CHECK: %__promise = alloca %"struct.std::experimental::coroutine_traits<int>::promise_type"
+  // CHECK: %call = call i8* @_Znwm(
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJiEE12promise_type12return_valueEi(%"struct.std::experimental::coroutine_traits<int>::promise_type"* %__promise, i32 42)
+  // CHECK: call void @_ZdlPv
+  co_return 42;
+}