Merge pull request swiftlang#36810 from nate-chandler/concurrency/irgen/rdar76372871

[IRGen] Restored single refcounted object opt for async partial applies.

Author: nate-chandler

lib/IRGen/GenFunc.cpp

@@ -1511,6 +1511,8 @@ static llvm::Value *emitPartialApplicationForwarder(IRGenModule &IGM,
     auto argIndex = emission->getCurrentArgumentIndex();
     if (haveContextArgument)
       argIndex += polyArgs.size();
+    if (origType->isAsync())
+      argIndex += 1;
 
     llvm::Type *expectedArgTy = origSig.getType()->getParamType(argIndex);
 
@@ -1858,13 +1860,6 @@ Optional<StackAddress> irgen::emitFunctionPartialApplication(
   auto bindings = NecessaryBindings::forPartialApplyForwarder(
       IGF.IGM, origType, subs, considerParameterSources);
 
-  // TODO: Revisit.  Now that async thick functions are always represented with
-  //       an async function pointer, it should again be possible to allow
-  //       contexts that consist of only a single swift refcounted value.
-  if (origType->isAsync()) {
-    hasSingleSwiftRefcountedContext = No;
-  }
-
   if (!bindings.empty()) {
     hasSingleSwiftRefcountedContext = No;
     auto bindingsSize = bindings.getBufferSize(IGF.IGM);
@@ -2025,13 +2020,8 @@ Optional<StackAddress> irgen::emitFunctionPartialApplication(
       emitPartialApplicationForwarder(IGF.IGM, staticFn, fnContext != nullptr,
                                       origSig, origType, substType,
                                       outType, subs, nullptr, argConventions);
-    if (origType->isAsync()) {
-      llvm_unreachable(
-          "async functions never have a single refcounted context");
-    } else {
-      forwarder = emitPointerAuthSign(IGF, forwarder, outAuthInfo);
-      forwarder = IGF.Builder.CreateBitCast(forwarder, IGF.IGM.Int8PtrTy);
-    }
+    forwarder = emitPointerAuthSign(IGF, forwarder, outAuthInfo);
+    forwarder = IGF.Builder.CreateBitCast(forwarder, IGF.IGM.Int8PtrTy);
     out.add(forwarder);
 
     llvm::Value *ctx = args.claimNext();

lib/IRGen/IRGenSIL.cpp

@@ -3215,9 +3215,6 @@ static bool isSimplePartialApply(IRGenFunction &IGF, PartialApplyInst *i) {
   // The callee type must use the `method` convention.
   auto calleeTy = i->getCallee()->getType().castTo<SILFunctionType>();
   auto resultTy = i->getFunctionType();
-
-  if (calleeTy->isAsync())
-    return false;
 
   if (calleeTy->getRepresentation() != SILFunctionTypeRepresentation::Method)
     return false;

test/IRGen/async/partial_apply.sil

@@ -617,3 +617,13 @@ bb0(%0 : $Int):
   %3 = partial_apply [callee_guaranteed] %2(%0) : $@convention(thin) @async (Int, Int) -> (Int, @error Error)
   return %3 : $@async @callee_guaranteed (Int) -> (Int, @error Error)
 }
+
+sil @method_with_out_return : $@convention(method) @async (SwiftClass) -> @out Int
+
+sil @partially_apply_method_with_out_return : $@convention(thin) (SwiftClass) -> @async @callee_guaranteed () -> @out Int {
+bb(%class : $SwiftClass):
+  %method_with_out_param = function_ref @method_with_out_return : $@convention(method) @async (SwiftClass) -> @out Int
+  %result = partial_apply [callee_guaranteed] %method_with_out_param(%class) : $@convention(method) @async (SwiftClass) -> @out Int
+  return %result : $@callee_guaranteed @async () -> @out Int
+}
+

test/IRGen/async/partial_apply_forwarder.sil

@@ -29,7 +29,6 @@ bb0(%0 : $*S):
   return %2 : $@async @callee_guaranteed (@in O) -> ()
 }
 
-// CHECK-LABEL: define internal swift{{(tail)?}}cc void @"$s23unspecialized_uncurriedTA"(
 sil hidden @specialized_curried : $@async @convention(thin) (@owned E) -> @owned @async @callee_guaranteed () -> @owned D<C> {
 bb0(%0 : $E):
   %1 = function_ref @unspecialized_uncurried : $@async @convention(method) <τ_0_0 where τ_0_0 : P> (@guaranteed E) -> @owned D<τ_0_0>