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

As part of bringup, specifically in order to support storing the size of
the async context as the first entry in the thick context, the
optimization that allows the partial application of a single refcounted
object to avoid the allocation of a thick context was disabled.

Now that we have async function pointers for partial application
forwarders, that rationale is moot, so, here, the optimization is
restored.

rdar://76372871

Author: nate-chandler

lib/IRGen/GenFunc.cpp

@@ -1506,6 +1506,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);
 
@@ -1853,13 +1855,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);
@@ -2020,13 +2015,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>