[coroutines] Add emission of initial and final suspends

https://reviews.llvm.org/D31608

llvm-svn: 303603

3 files changed, 58 insertions, 11 deletions

diff --git a/clang/lib/CodeGen/CGCoroutine.cpp b/clang/lib/CodeGen/CGCoroutine.cpp
index 26291c9978b..bb98c248908 100644
--- a/clang/lib/CodeGen/CGCoroutine.cpp
+++ b/clang/lib/CodeGen/CGCoroutine.cpp
@@ -388,7 +388,10 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
 
     CurCoro.Data->FinalJD = getJumpDestInCurrentScope(FinalBB);
 
-    // FIXME: Emit initial suspend and more before the body.
+    // FIXME: Emit param moves.
+
+    CurCoro.Data->CurrentAwaitKind = AwaitKind::Init;
+    EmitStmt(S.getInitSuspendStmt());
 
     CurCoro.Data->CurrentAwaitKind = AwaitKind::Normal;
 
@@ -410,7 +413,8 @@ void CodeGenFunction::EmitCoroutineBody(const CoroutineBodyStmt &S) {
     const bool HasCoreturns = CurCoro.Data->CoreturnCount > 0;
     if (CanFallthrough || HasCoreturns) {
       EmitBlock(FinalBB);
-      // FIXME: Emit final suspend.
+      CurCoro.Data->CurrentAwaitKind = AwaitKind::Final;
+      EmitStmt(S.getFinalSuspendStmt());
     }
   }
 
diff --git a/clang/test/CodeGenCoroutines/coro-await.cpp b/clang/test/CodeGenCoroutines/coro-await.cpp
index 8a9b9919e36..dfeeba035f6 100644
--- a/clang/test/CodeGenCoroutines/coro-await.cpp
+++ b/clang/test/CodeGenCoroutines/coro-await.cpp
@@ -21,6 +21,17 @@ struct coroutine_handle : coroutine_handle<> {
 }
 }
 
+struct init_susp {
+  bool await_ready();
+  void await_suspend(std::experimental::coroutine_handle<>);
+  void await_resume();
+};
+struct final_susp {
+  bool await_ready();
+  void await_suspend(std::experimental::coroutine_handle<>);
+  void await_resume();
+};
+
 struct suspend_always {
   int stuff;
   bool await_ready();
@@ -32,8 +43,8 @@ template<>
 struct std::experimental::coroutine_traits<void> {
   struct promise_type {
     void get_return_object();
-    suspend_always initial_suspend();
-    suspend_always final_suspend();
+    init_susp initial_suspend();
+    final_susp final_suspend();
     void return_void();
   };
 };
@@ -42,6 +53,13 @@ struct std::experimental::coroutine_traits<void> {
 extern "C" void f0() {
   // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(
 
+  // See if initial_suspend was issued:
+  // ----------------------------------
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_type15initial_suspendEv(
+  // CHECK-NEXT: call zeroext i1 @_ZN9init_susp11await_readyEv(%struct.init_susp*
+  // CHECK: %[[INITSP_ID:.+]] = call token @llvm.coro.save(
+  // CHECK: call i8 @llvm.coro.suspend(token %[[INITSP_ID]], i1 false)
+
   co_await suspend_always{};
   // See if we need to suspend:
   // --------------------------
@@ -76,6 +94,13 @@ extern "C" void f0() {
   // --------------------------
   // CHECK: [[READY_BB]]:
   // CHECK:  call void @_ZN14suspend_always12await_resumeEv(%struct.suspend_always* %[[AWAITABLE]])
+
+  // See if final_suspend was issued:
+  // ----------------------------------
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_type13final_suspendEv(
+  // CHECK-NEXT: call zeroext i1 @_ZN10final_susp11await_readyEv(%struct.final_susp*
+  // CHECK: %[[FINALSP_ID:.+]] = call token @llvm.coro.save(
+  // CHECK: call i8 @llvm.coro.suspend(token %[[FINALSP_ID]], i1 true)
 }
 
 struct suspend_maybe {
@@ -91,8 +116,8 @@ template<>
 struct std::experimental::coroutine_traits<void,int> {
   struct promise_type {
     void get_return_object();
-    suspend_always initial_suspend();
-    suspend_always final_suspend();
+    init_susp initial_suspend();
+    final_susp final_suspend();
     void return_void();
     suspend_maybe yield_value(int);
   };
@@ -228,3 +253,21 @@ extern "C" void TestOpAwait() {
   // CHECK: call void @_ZN5MyAggawEv(%struct.MyAgg* %
   // CHECK: call void @_ZN11AggrAwaiter12await_resumeEv(%struct.Aggr* sret %
 }
+
+// CHECK-LABEL: EndlessLoop(
+extern "C" void EndlessLoop() {
+  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(
+
+  // See if initial_suspend was issued:
+  // ----------------------------------
+  // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_type15initial_suspendEv(
+  // CHECK-NEXT: call zeroext i1 @_ZN9init_susp11await_readyEv(%struct.init_susp*
+
+  for (;;)
+    co_await suspend_always{};
+
+  // Verify that final_suspend was NOT issued:
+  // ----------------------------------
+  // CHECK-NOT: call void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_type13final_suspendEv(
+  // CHECK-NOT: call zeroext i1 @_ZN10final_susp11await_readyEv(%struct.final_susp*
+}
diff --git a/clang/test/CodeGenCoroutines/coro-cleanup.cpp b/clang/test/CodeGenCoroutines/coro-cleanup.cpp
index 519282e2b6b..ba1246979e5 100644
--- a/clang/test/CodeGenCoroutines/coro-cleanup.cpp
+++ b/clang/test/CodeGenCoroutines/coro-cleanup.cpp
@@ -46,6 +46,11 @@ void f() {
   // CHECK: invoke void @_ZNSt12experimental16coroutine_traitsIJvEE12promise_typeC1Ev(
   // CHECK-NEXT: to label %{{.+}} unwind label %[[DeallocPad:.+]]
 
+  // CHECK: [[DeallocPad]]:
+  // CHECK-NEXT: landingpad
+  // CHECK-NEXT:   cleanup
+  // CHECK: br label %[[Dealloc:.+]]
+
   Cleanup cleanup;
   may_throw();
 
@@ -54,11 +59,6 @@ void f() {
   // CHECK: invoke void @_Z9may_throwv(
   // CHECK-NEXT: to label %{{.+}} unwind label %[[CatchPad:.+]]
 
-  // CHECK: [[DeallocPad]]:
-  // CHECK-NEXT: landingpad
-  // CHECK-NEXT:   cleanup
-  // CHECK: br label %[[Dealloc:.+]]
-
   // CHECK: [[CatchPad]]:
   // CHECK-NEXT:  landingpad
   // CHECK-NEXT:       catch i8* null