Correctly forward the implicit nonisolated(nonsending) parameter in thunks (swiftlang#82705)

Also, I discovered that we don't apply nonisolated(nonsending) to
function types in the new mode. That's one for a different patch.

Fixes rdar://154401813 and rdar://154137740

Author: rjmccall

lib/SILGen/SILGenAvailability.cpp

@@ -508,7 +508,9 @@ void SILGenFunction::emitBackDeploymentThunk(SILDeclRef thunk) {
   SmallVector<ManagedValue, 4> params;
   SmallVector<ManagedValue, 4> indirectParams;
   SmallVector<ManagedValue, 4> indirectErrorResults;
-  collectThunkParams(loc, params, &indirectParams, &indirectErrorResults);
+  ManagedValue implicitIsolationParam;
+  collectThunkParams(loc, params, &indirectParams, &indirectErrorResults,
+                     &implicitIsolationParam);
 
   // Build up the list of arguments that we're going to invoke the real
   // function with.
@@ -519,6 +521,9 @@ void SILGenFunction::emitBackDeploymentThunk(SILDeclRef thunk) {
   for (auto indirectErrorResult : indirectErrorResults) {
     paramsForForwarding.emplace_back(indirectErrorResult.getLValueAddress());
   }
+  if (implicitIsolationParam) {
+    paramsForForwarding.emplace_back(implicitIsolationParam.forward(*this));
+  }
 
   for (auto param : params) {
     // We're going to directly call either the original function or the fallback

lib/SILGen/SILGenFunction.h

@@ -2623,7 +2623,8 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   void collectThunkParams(
       SILLocation loc, SmallVectorImpl<ManagedValue> &params,
       SmallVectorImpl<ManagedValue> *indirectResultParams = nullptr,
-      SmallVectorImpl<ManagedValue> *indirectErrorParams = nullptr);
+      SmallVectorImpl<ManagedValue> *indirectErrorParams = nullptr,
+      ManagedValue *implicitIsolationParam = nullptr);
 
   /// Build the type of a function transformation thunk.
   CanSILFunctionType buildThunkType(CanSILFunctionType &sourceType,

lib/SILGen/SILGenPoly.cpp

@@ -987,7 +987,8 @@ ManagedValue Transform::transformTuple(ManagedValue inputTuple,
 void SILGenFunction::collectThunkParams(
     SILLocation loc, SmallVectorImpl<ManagedValue> &params,
     SmallVectorImpl<ManagedValue> *indirectResults,
-    SmallVectorImpl<ManagedValue> *indirectErrors) {
+    SmallVectorImpl<ManagedValue> *indirectErrors,
+    ManagedValue *implicitIsolation) {
   // Add the indirect results.
   for (auto resultTy : F.getConventions().getIndirectSILResultTypes(
            getTypeExpansionContext())) {
@@ -1029,8 +1030,12 @@ void SILGenFunction::collectThunkParams(
     // If our thunk has an implicit param and we are being asked to forward it,
     // to the callee, skip it. We are going to handle it especially later.
     if (param.hasOption(SILParameterInfo::ImplicitLeading) &&
-        param.hasOption(SILParameterInfo::Isolated))
+        param.hasOption(SILParameterInfo::Isolated)) {
+      if (implicitIsolation)
+        *implicitIsolation = functionArgument;
       continue;
+    }
+
     params.push_back(functionArgument);
   }
 }
@@ -5463,9 +5468,18 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
       options |= ThunkGenFlag::CalleeHasImplicitIsolatedParam;
   }
 
+  // Collect the thunk parameters. We don't need to collect the indirect
+  // error parameter because it'll be stored as IndirectErrorResult, which
+  // gets used implicitly by emitApplyWithRethrow.
   SmallVector<ManagedValue, 8> params;
   SmallVector<ManagedValue, 4> indirectResultParams;
-  SGF.collectThunkParams(loc, params, &indirectResultParams);
+  ManagedValue implicitIsolationParam;
+  SGF.collectThunkParams(loc, params, &indirectResultParams,
+                         /*indirect error params*/ nullptr,
+                         &implicitIsolationParam);
+
+  assert(bool(options & ThunkGenFlag::ThunkHasImplicitIsolatedParam)
+           == implicitIsolationParam.isValid());
 
   // Ignore the self parameter at the SIL level. IRGen will use it to
   // recover type metadata.
@@ -5511,9 +5525,11 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
     case FunctionTypeIsolation::Kind::NonIsolated:
       break;
 
+    // For a function for caller isolation, we'll have to figure out what the
+    // output function's formal isolation is. This is quite doable, but we don't
+    // have to do it yet.
     case FunctionTypeIsolation::Kind::NonIsolatedCaller:
-      hopToIsolatedParameter = true;
-      break;
+      llvm_unreachable("synchronous function has caller isolation?");
 
     // For a function with parameter isolation, we'll have to dig the
     // argument out after translation but before making the call.
@@ -5595,12 +5611,15 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
                               /*mandatory*/false);
   }
 
-  // If we are thunking a nonisolated caller to nonisolated or global actor, we
-  // need to load the actor.
+  // If the input function has caller isolation, we need to fill in that
+  // argument with the formal isolation of the output function.
   if (options.contains(ThunkGenFlag::CalleeHasImplicitIsolatedParam)) {
     auto outputIsolation = outputSubstType->getIsolation();
     switch (outputIsolation.getKind()) {
     case FunctionTypeIsolation::Kind::NonIsolated:
+    case FunctionTypeIsolation::Kind::Erased:
+      // Converting a caller-isolated function to @isolated(any) makes
+      // it @concurrent. In either case, emit a nil actor reference.
       argValues.push_back(SGF.emitNonIsolatedIsolation(loc).getValue());
       break;
     case FunctionTypeIsolation::Kind::GlobalActor: {
@@ -5610,11 +5629,17 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
           SGF.emitGlobalActorIsolation(loc, globalActor).getValue());
       break;
     }
+    case FunctionTypeIsolation::Kind::NonIsolatedCaller: {
+      argValues.push_back(implicitIsolationParam.getValue());
+      break;
+    }
     case FunctionTypeIsolation::Kind::Parameter:
-    case FunctionTypeIsolation::Kind::Erased:
-    case FunctionTypeIsolation::Kind::NonIsolatedCaller:
+      // This would require a conversion from a type with caller
+      // isolation to a type with parameter isolation, which is not
+      // currently allowed and probably won't ever be. Anyway, to
+      // implement it, we'd need to borrow the isolated parameter
+      // and wrap it up as an `Optional<any Actor>`.
       llvm_unreachable("Should never see this");
-      break;
     }
   }
 
@@ -5965,7 +5990,9 @@ static void buildWithoutActuallyEscapingThunkBody(SILGenFunction &SGF,
   SmallVector<ManagedValue, 8> params;
   SmallVector<ManagedValue, 8> indirectResults;
   SmallVector<ManagedValue, 1> indirectErrorResults;
-  SGF.collectThunkParams(loc, params, &indirectResults, &indirectErrorResults);
+  ManagedValue implicitIsolationParam;
+  SGF.collectThunkParams(loc, params, &indirectResults, &indirectErrorResults,
+                         &implicitIsolationParam);
 
   // Ignore the self parameter at the SIL level. IRGen will use it to
   // recover type metadata.
@@ -5986,6 +6013,11 @@ static void buildWithoutActuallyEscapingThunkBody(SILGenFunction &SGF,
   for (auto indirectError : indirectErrorResults)
     argValues.push_back(indirectError.getLValueAddress());
 
+  // Forward the implicit isolation parameter.
+  if (implicitIsolationParam.isValid()) {
+    argValues.push_back(implicitIsolationParam.getValue());
+  }
+
    // Add the rest of the arguments.
   forwardFunctionArguments(SGF, loc, fnType, params, argValues);
 
@@ -6175,6 +6207,7 @@ ManagedValue SILGenFunction::getThunkedAutoDiffLinearMap(
     return getThunkedResult();
   thunk->setGenericEnvironment(genericEnv);
 
+  // FIXME: handle implicit isolation parameter here
   SILGenFunction thunkSGF(SGM, *thunk, FunctionDC);
   SmallVector<ManagedValue, 4> params;
   SmallVector<ManagedValue, 4> thunkIndirectResults;
@@ -6523,6 +6556,7 @@ SILFunction *SILGenModule::getOrCreateCustomDerivativeThunk(
   thunk->setGenericEnvironment(thunkGenericEnv);
 
   SILGenFunction thunkSGF(*this, *thunk, customDerivativeFn->getDeclContext());
+  // FIXME: handle implicit isolation parameter here
   SmallVector<ManagedValue, 4> params;
   SmallVector<ManagedValue, 4> indirectResults;
   SmallVector<ManagedValue, 1> indirectErrorResults;

lib/SILGen/SILGenType.cpp

@@ -1298,8 +1298,9 @@ SILFunction *SILGenModule::emitDefaultOverride(SILDeclRef replacement,
   SmallVector<ManagedValue, 4> params;
   SmallVector<ManagedValue, 4> indirectResults;
   SmallVector<ManagedValue, 4> indirectErrors;
+  ManagedValue implicitIsolationParam;
   SGF.collectThunkParams(replacement.getDecl(), params, &indirectResults,
-                         &indirectErrors);
+                         &indirectErrors, &implicitIsolationParam);
 
   auto self = params.back();
 
@@ -1315,6 +1316,11 @@ SILFunction *SILGenModule::emitDefaultOverride(SILDeclRef replacement,
   for (auto result : indirectResults) {
     args.push_back(result.forward(SGF));
   }
+  // Indirect errors would go here, but we don't currently support
+  // throwing coroutines.
+  if (implicitIsolationParam.isValid()) {
+    args.push_back(implicitIsolationParam.forward(SGF));
+  }
   for (auto param : params) {
     args.push_back(param.forward(SGF));
   }

test/Concurrency/nonisolated_nonsending.swift

@@ -0,0 +1,64 @@
+// RUN: %target-swift-frontend -emit-sil -parse-as-library -swift-version 6 -enable-upcoming-feature NonisolatedNonsendingByDefault %s | %FileCheck %s
+
+// REQUIRES: swift_feature_NonisolatedNonsendingByDefault
+
+protocol P: Sendable {
+  func f() async
+}
+
+// Our map closure (abstracted as returning T) returns a function
+// value, which requires it to be reabstracted to the most general
+// abstraction pattern, i.e. from
+//     nonisolated(nonsending) () async -> ()
+// to
+//     nonisolated(nonsending) () async -> U, where U == ()
+// Note that we preserve nonisolated(nonsending).
+//
+// The thunk code did not expect the output function type to be
+// nonisolated(nonsending), so it didn't handle and propagate the
+// implicit isolation argument correctly.
+func testPartialApplication(p: [any P]) async {
+  _ = p.map { $0.f }
+}
+// CHECK-LABEL: sil private @$s22nonisolated_nonsending22testPartialApplication1pySayAA1P_pG_tYaFyyYaYbYCcAaD_pXEfU_ :
+// CHECK: function_ref @$sScA_pSgIeghHgIL_AAytIeghHgILr_TR :
+
+//   Reabstraction thunk from caller-isolated () -> () to caller-isolated () -> T
+// CHECK-LABEL: sil shared [transparent] [reabstraction_thunk] @$sScA_pSgIeghHgIL_AAytIeghHgILr_TR :
+// CHECK:       bb0(%0 : $*(), %1 : $Optional<any Actor>, %2 : $@Sendable @async @callee_guaranteed (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>) -> ()):
+// CHECK-NEXT:    apply %2(%1) :
+
+func takesGenericAsyncFunction<T>(_ fn: nonisolated(nonsending) (T) async -> Void) {}
+
+// This is a more direct test of the partial-application code above.
+// Note that we have to make the functions explicitly nonisolated(nonsending)
+// because NonisolatedNonsendingByDefault only applies to declarations,
+// not function types in the abstract.
+func testReabstractionPreservingCallerIsolation(fn: nonisolated(nonsending) (Int) async -> Void) {
+  takesGenericAsyncFunction(fn)
+}
+// CHECK-LABEL: sil hidden @$s22nonisolated_nonsending42testReabstractionPreservingCallerIsolation2fnyySiYaYCXE_tF :
+// CHECK:       bb0(%0 : $@noescape @async @callee_guaranteed (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>, Int) -> ()):
+// CHECK:         [[THUNK:%.*]] = function_ref @$sScA_pSgSiIgHgILy_AASiIegHgILn_TR :
+
+// CHECK-LABEL: sil shared [transparent] [reabstraction_thunk] @$sScA_pSgSiIgHgILy_AASiIegHgILn_TR :
+// CHECK:       bb0(%0 : $Optional<any Actor>, %1 : $*Int, %2 : $@noescape @async @callee_guaranteed (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>, Int) -> ()):
+// CHECK-NEXT:    %3 = load %1
+// CHECK-NEXT:    apply %2(%0, %3)
+
+func takesAsyncIsolatedAnyFunction(_ fn: @isolated(any) () async -> Void) {}
+func takesGenericAsyncIsolatedAnyFunction<T>(_ fn: @isolated(any) (T) async -> Void) {}
+
+// These would be good to test, but we apparently reject this conversion.
+#if false
+
+// The same bug, but converting to an @isolated(any) function type.
+func testConversionToIsolatedAny(fn: nonisolated(nonsending) () async -> Void) {
+  takesAsyncIsolatedAnyFunction(fn)
+}
+
+func testReabstractionToIsolatedAny(fn: nonisolated(nonsending) (Int) async -> Void) {
+  takesGenericAsyncIsolatedAnyFunction(fn)
+}
+
+#endif

test/SILGen/back_deployed_attr.swift

@@ -85,3 +85,18 @@ public func backDeployedCaller(_ s: inout S<Z>) {
   // CHECK: function_ref @$s11back_deploy1SV1xxvsTwb : $@convention(method) <τ_0_0> (@in τ_0_0, @inout S<τ_0_0>) -> ()
   s.x = Z()
 }
+
+// The same bug from test/Concurrency/nonisolated_nonsending.swift also applied to
+// back-deployment thunks.
+
+// CHECK-LABEL: sil non_abi [serialized] [back_deployed_thunk] [ossa] @$s11back_deploy0A29DeployedNonisolatedNonsendingSiyYaFTwb :
+@backDeployed(before: macOS 10.52)
+nonisolated(nonsending)
+public func backDeployedNonisolatedNonsending() async -> Int {
+  // CHECK: bb0(%0 : @guaranteed $Optional<any Actor>):
+  // CHECK:   [[FALLBACK_FN:%.*]] = function_ref @$s11back_deploy0A29DeployedNonisolatedNonsendingSiyYaFTwB :
+  // CHECK:   apply [[FALLBACK_FN]](%0)
+  // CHECK:   [[SHIPPING_FN:%.*]] = function_ref @$s11back_deploy0A29DeployedNonisolatedNonsendingSiyYaF :
+  // CHECK:   apply [[SHIPPING_FN]](%0)
+  return 0
+}