1 files changed, 33 insertions, 7 deletions
diff --git a/test/CodeGenCoroutines/coro-alloc.cpp b/test/CodeGenCoroutines/coro-alloc.cpp
index 3a65736f6c03a..820201db357fd 100644
--- a/test/CodeGenCoroutines/coro-alloc.cpp
+++ b/test/CodeGenCoroutines/coro-alloc.cpp
@@ -56,12 +56,28 @@ struct std::experimental::coroutine_traits<void, global_new_delete_tag> {
 // CHECK-LABEL: f0(
 extern "C" void f0(global_new_delete_tag) {
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
+  // CHECK: %[[NeedAlloc:.+]] = call i1 @llvm.coro.alloc(token %[[ID]])
+  // CHECK: br i1 %[[NeedAlloc]], label %[[AllocBB:.+]], label %[[InitBB:.+]]
+
+  // CHECK: [[AllocBB]]:
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
-  // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
+  // CHECK: %[[MEM:.+]] = call i8* @_Znwm(i64 %[[SIZE]])
+  // CHECK: br label %[[InitBB]]
+
+  // CHECK: [[InitBB]]:
+  // CHECK: %[[PHI:.+]] = phi i8* [ null, %{{.+}} ], [ %call, %[[AllocBB]] ]
+  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(token %[[ID]], i8* %[[PHI]])
 
-  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
+  // CHECK: %[[NeedDealloc:.+]] = icmp ne i8* %[[MEM]], null
+  // CHECK: br i1 %[[NeedDealloc]], label %[[FreeBB:.+]], label %[[Afterwards:.+]]
+
+  // CHECK: [[FreeBB]]:
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
+  // CHECK: br label %[[Afterwards]]
+
+  // CHECK: [[Afterwards]]:
+  // CHECK: ret void
   co_return;
 }
 
@@ -84,7 +100,7 @@ extern "C" void f1(promise_new_tag ) {
   // 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: %[[FRAME:.+]] = call i8* @llvm.coro.begin(
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZdlPv(i8* %[[MEM]])
   co_return;
@@ -109,7 +125,7 @@ extern "C" void f2(promise_delete_tag) {
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
-  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
+  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(
   // CHECK: %[[MEM:.+]] = call i8* @llvm.coro.free(token %[[ID]], i8* %[[FRAME]])
   // CHECK: call void @_ZNSt12experimental16coroutine_traitsIJv18promise_delete_tagEE12promise_typedlEPv(i8* %[[MEM]])
   co_return;
@@ -134,7 +150,7 @@ extern "C" void f3(promise_sized_delete_tag) {
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: call i8* @_Znwm(i64 %[[SIZE]])
 
-  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.frame()
+  // CHECK: %[[FRAME:.+]] = call i8* @llvm.coro.begin(
   // 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]])
@@ -156,6 +172,7 @@ struct std::experimental::coroutine_traits<int, promise_on_alloc_failure_tag> {
 
 // CHECK-LABEL: f4(
 extern "C" int f4(promise_on_alloc_failure_tag) {
+  // CHECK: %[[RetVal:.+]] = alloca i32
   // CHECK: %[[ID:.+]] = call token @llvm.coro.id(i32 16
   // CHECK: %[[SIZE:.+]] = call i64 @llvm.coro.size.i64()
   // CHECK: %[[MEM:.+]] = call i8* @_ZnwmRKSt9nothrow_t(i64 %[[SIZE]], %"struct.std::nothrow_t"* dereferenceable(1) @_ZStL7nothrow)
@@ -163,7 +180,16 @@ extern "C" int f4(promise_on_alloc_failure_tag) {
   // CHECK: br i1 %[[OK]], label %[[OKBB:.+]], label %[[ERRBB:.+]]
 
   // CHECK: [[ERRBB]]:
-  // CHECK: %[[RETVAL:.+]] = call i32 @_ZNSt12experimental16coroutine_traitsIJi28promise_on_alloc_failure_tagEE12promise_type39get_return_object_on_allocation_failureEv(
-  // CHECK: ret i32 %[[RETVAL]]
+  // CHECK:   %[[FailRet:.+]] = call i32 @_ZNSt12experimental16coroutine_traitsIJi28promise_on_alloc_failure_tagEE12promise_type39get_return_object_on_allocation_failureEv(
+  // CHECK:   store i32 %[[FailRet]], i32* %[[RetVal]]
+  // CHECK:   br label %[[RetBB:.+]]
+
+  // CHECK: [[OKBB]]:
+  // CHECK:   %[[OkRet:.+]] = call i32 @_ZNSt12experimental16coroutine_traitsIJi28promise_on_alloc_failure_tagEE12promise_type17get_return_objectEv(
+  // CHECK:   store i32 %[[OkRet]], i32* %[[RetVal]]
+
+  // CHECK: [[RetBB]]:
+  // CHECK:   %[[LoadRet:.+]] = load i32, i32* %[[RetVal]], align 4
+  // CHECK:   ret i32 %[[LoadRet]]
   co_return;
 }