[IRGen] Support indirect results for coroutines (#75322)

Author: asl

lib/IRGen/GenCall.cpp

@@ -760,9 +760,6 @@ void SignatureExpansion::expandCoroutineResult(bool forContinuation) {
 
     // Yield-once coroutines may optionaly return a value from the continuation.
     case SILCoroutineKind::YieldOnce: {
-      auto fnConv = getSILFuncConventions();
-
-      assert(fnConv.getNumIndirectSILResults() == 0);
       // Ensure that no parameters were added before to correctly record their ABI
       // details.
       assert(ParamIRTypes.empty());
@@ -1930,6 +1927,17 @@ void SignatureExpansion::expandParameters(
   case SILCoroutineKind::YieldOnce:
   case SILCoroutineKind::YieldMany:
     addCoroutineContextParameter();
+
+    // Add indirect results as parameters. Similar to
+    // expandIndirectResults, but it doesn't add sret attribute,
+    // because the function has direct results (a continuation pointer
+    // and yield results).
+    auto fnConv = getSILFuncConventions();
+    for (auto indirectResultType : fnConv.getIndirectSILResultTypes(
+           IGM.getMaximalTypeExpansionContext())) {
+      auto storageTy = IGM.getStorageType(indirectResultType);
+      addPointerParameter(storageTy);
+    }
     break;
   }
 
@@ -6330,13 +6338,12 @@ void irgen::emitYieldOnceCoroutineResult(IRGenFunction &IGF, Explosion &result,
       // Capture results via result token
       resultToken =
         Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end_results, coroResults);
-
-        Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end,
-                              {handle,
-                               /*is unwind*/ Builder.getFalse(),
-                               resultToken});
-        Builder.CreateUnreachable();
     }
+    Builder.CreateIntrinsicCall(llvm::Intrinsic::coro_end,
+                                {handle,
+                                 /*is unwind*/ Builder.getFalse(),
+                                 resultToken});
+    Builder.CreateUnreachable();
   } else {
     if (coroResults.empty()) {
       // No results, we do not need to change anything around existing coro.end

lib/IRGen/IRGenSIL.cpp

@@ -2114,6 +2114,19 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
       getNativeCCEntryPointArgumentEmission(IGF, *entry, allParamValues);
   auto funcTy = IGF.CurSILFn->getLoweredFunctionType();
 
+  // Coroutine context should be the first parameter.
+  // Indirect returns (if present) follow it.
+  switch (funcTy->getCoroutineKind()) {
+  case SILCoroutineKind::None:
+    break;
+  case SILCoroutineKind::YieldOnce:
+    emitYieldOnceCoroutineEntry(IGF, funcTy, *emission);
+    break;
+  case SILCoroutineKind::YieldMany:
+    emitYieldManyCoroutineEntry(IGF, funcTy, *emission);
+    break;
+  }
+  
   // Map the indirect return if present.
   ArrayRef<SILArgument *> params = emitEntryPointIndirectReturn(
       *emission, IGF, entry, funcTy, [&](SILType retType) -> bool {
@@ -2130,18 +2143,6 @@ static void emitEntryPointArgumentsNativeCC(IRGenSILFunction &IGF,
                                    witnessMetadata);
   }
 
-  // The coroutine context should be the first parameter.
-  switch (funcTy->getCoroutineKind()) {
-  case SILCoroutineKind::None:
-    break;
-  case SILCoroutineKind::YieldOnce:
-    emitYieldOnceCoroutineEntry(IGF, funcTy, *emission);
-    break;
-  case SILCoroutineKind::YieldMany:
-    emitYieldManyCoroutineEntry(IGF, funcTy, *emission);
-    break;
-  }
-
   SILFunctionConventions fnConv(funcTy, IGF.getSILModule());
   if (funcTy->isAsync()) {
     emitAsyncFunctionEntry(IGF, getAsyncContextLayout(IGF.IGM, IGF.CurSILFn),

test/IRGen/yield_result.sil

@@ -91,3 +91,60 @@ entry:
   %retf = tuple (%ret : $Builtin.Int64, %ret2_1 : $Builtin.Int64, %ret2_2 : $Builtin.Int64)
   return %retf : $(Builtin.Int64, Builtin.Int64, Builtin.Int64)
 }
+
+// CHECK-LABEL: coro_ret_indirect
+// CHECK-SAME: ptr{{.*}} [[CTX:%.*]], ptr{{.*}} [[INDIRECT_RET:%.*]], ptr{{.*}} [[ARG:%.*]], ptr{{.*}} [[TYPE:%.*]])
+sil @coro_ret_indirect : $@yield_once @convention(thin) <T> (@in T) -> (@yields @in T, @out T) {
+bb0(%outt : $*T, %t : $*T):
+  // CHECK-32:      [[ID:%.*]] = call token @llvm.coro.id.retcon.once(i32 [[BUFFER_SIZE]], i32 [[BUFFER_ALIGN:4]], ptr [[CTX]], ptr @"$sxxxlIetAirYi_TC{{(.ptrauth)?}}", ptr @malloc, ptr @free)
+  // CHECK-64:      [[ID:%.*]] = call token @llvm.coro.id.retcon.once(i32 [[BUFFER_SIZE]], i32 [[BUFFER_ALIGN:8]], ptr [[CTX]], ptr @"$sxxxlIetAirYi_TC{{(.ptrauth)?}}", ptr @malloc, ptr @free)
+
+  // CHECK:    [[IS_UNWIND:%.*]] = call i1 (...) @llvm.coro.suspend.retcon.i1(ptr [[ARG]])
+  // CHECK:    br i1 [[IS_UNWIND]], label %[[UNWIND_BB:.*]], label %[[RESUME_BB:.*]]
+
+  // CHECK:[[RESUME_BB]]:
+  // CHECK: [[VW_PTR:%.*]] = getelementptr inbounds ptr, ptr [[TYPE]], i64 -1
+  // CHECK: [[VW:%.*]] = load ptr, ptr [[VW_PTR]]
+  // CHECK: [[ASSIGN_PTR:%.*]] = getelementptr inbounds ptr, ptr [[VW]], i32 3
+  // CHECK: [[ASSIGN:%.*]] = load ptr, ptr [[ASSIGN_PTR]]
+  // CHECK: call ptr [[ASSIGN]](ptr [[INDIRECT_RET]], ptr [[ARG]], ptr [[TYPE]]) #2
+
+  yield (%t : $*T), resume bb1, unwind bb2
+
+bb1:
+  copy_addr %t to %outt : $*T
+  %r = tuple ()
+  return %r : $()
+
+bb2:
+  unwind
+}
+
+// CHECK-LABEL: @test_coro_ret_indirect
+// CHECK-SAME: (i64 [[ARG:%.*]])
+sil [ossa] @test_coro_ret_indirect : $(Builtin.Int64) -> () {
+bb0(%0 : $Builtin.Int64):
+  // CHECK: [[ARG_COPY:%.*]] = alloca i64
+  // CHECK: [[INDIRECT_RET:%.*]] = alloca i64
+  // CHECK: [[FRAME:%.*]] = alloca [32 x i8]
+  %coro = function_ref @coro_ret_indirect : $@yield_once @convention(thin) <T> (@in T) -> (@yields @in T, @out T)
+  %temp = alloc_stack $Builtin.Int64
+  store %0 to [trivial] %temp : $*Builtin.Int64
+
+  %out = alloc_stack $Builtin.Int64
+
+  // CHECK: store i64 [[ARG]], ptr [[ARG_COPY]]
+  // CHECK: [[CTX:%.*]] = getelementptr inbounds [32 x i8], ptr [[FRAME]], i32 0, i32 0
+  // CHECK: [[CORO:%.*]] = call ptr @llvm.coro.prepare.retcon(ptr @coro_ret_indirect)
+  // CHECK: [[FRAME:%.*]] = call swiftcc { ptr, ptr } [[CORO]](ptr{{.*}} [[CTX]], ptr [[INDIRECT_RET]], ptr noalias [[ARG_COPY]], ptr getelementptr inbounds (%swift.full_existential_type, ptr @{{.*}}
+
+  (%f1, %token) = begin_apply %coro<Builtin.Int64>(%out, %temp) : $@yield_once @convention(thin) <T> (@in T) -> (@yields @in T, @out T)
+
+  %f2 = end_apply %token as $()
+
+  dealloc_stack %out : $*Builtin.Int64
+  dealloc_stack %temp : $*Builtin.Int64
+
+  %r = tuple ()
+  return %r : $()
+}