[silgen] Make sure that thunks that convert to/from nonisolated(nonsending) handle hopping correctly.

Specifically:

1. When we convert a function to nonisolated(nonsending), we need to
make sure that in the thunk we hop upon return since nonisolated(nonsending)
functions are assumed to preserve the caller's isolation.

2. When we convert a function from nonisolated(nonsending), we need to
make sure that in the thunk we hop onto the actor that we are passing in as the
isolated parameter of the nonisolated(nonsending) function. This ensures that
the nonisolated(nonsending) function can assume that it is already on its
isolated parameter's actor at function entry.

rdar://155905383

Author: gottesmm

lib/SILGen/SILGenConcurrency.cpp

@@ -680,6 +680,35 @@ SILGenFunction::emitHopToTargetActor(SILLocation loc,
   }
 }
 
+namespace {
+
+class HopToActorCleanup : public Cleanup {
+  SILValue value;
+
+public:
+  HopToActorCleanup(SILValue value) : value(value) {}
+
+  void emit(SILGenFunction &SGF, CleanupLocation l,
+            ForUnwind_t forUnwind) override {
+    SGF.B.createHopToExecutor(l, value, false /*mandatory*/);
+  }
+
+  void dump(SILGenFunction &) const override {
+#ifndef NDEBUG
+    llvm::errs() << "HopToExecutorCleanup\n"
+                 << "State:" << getState() << "\n"
+                 << "Value:" << value << "\n";
+#endif
+  }
+};
+} // end anonymous namespace
+
+CleanupHandle SILGenFunction::emitScopedHopToTargetActor(SILLocation loc,
+                                                         SILValue actor) {
+  Cleanups.pushCleanup<HopToActorCleanup>(actor);
+  return Cleanups.getTopCleanup();
+}
+
 ExecutorBreadcrumb SILGenFunction::emitHopToTargetExecutor(
     SILLocation loc, SILValue executor) {
   // Record that we need to hop back to the current executor.

lib/SILGen/SILGenFunction.h

@@ -1333,6 +1333,12 @@ class LLVM_LIBRARY_VISIBILITY SILGenFunction
   emitHopToTargetActor(SILLocation loc, std::optional<ActorIsolation> actorIso,
                        std::optional<ManagedValue> actorSelf);
 
+  /// Emit a cleanup for a hop back to a target actor.
+  ///
+  /// Used to ensure that along error paths and normal scope exit paths we hop
+  /// back automatically.
+  CleanupHandle emitScopedHopToTargetActor(SILLocation loc, SILValue actor);
+
   /// Emit a hop to the target executor, returning a breadcrumb with enough
   /// enough information to hop back.
   ///

lib/SILGen/SILGenPoly.cpp

@@ -5517,7 +5517,7 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
   // don't need to do anything.  But if the input is synchronous and the
   // executor is asynchronous, we need to treat this like any other call
   // to a synchronous function from an asynchronous context.
-  bool hopToIsolatedParameter = false;
+  bool hopToIsolatedParameterForSync = false;
   if (outputSubstType->isAsync() && !inputSubstType->isAsync()) {
     auto inputIsolation = inputSubstType->getIsolation();
     switch (inputIsolation.getKind()) {
@@ -5534,7 +5534,7 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
     // For a function with parameter isolation, we'll have to dig the
     // argument out after translation but before making the call.
     case FunctionTypeIsolation::Kind::Parameter:
-      hopToIsolatedParameter = true;
+      hopToIsolatedParameterForSync = true;
       break;
 
     // For a function with global-actor isolation, hop to the appropriate
@@ -5603,44 +5603,49 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
     argValues.push_back(*innerIndirectErrorAddr);
   }
 
-  // If we need to jump to an isolated parameter, do so before the call.
-  if (hopToIsolatedParameter) {
+  // If we need to jump to an isolated parameter for a synchronous function, do
+  // so before the call.
+  if (hopToIsolatedParameterForSync) {
     auto formalIsolatedIndex = getIsolatedParamIndex(inputSubstType);
     auto isolatedIndex = inputOrigType.getLoweredParamIndex(formalIsolatedIndex);
     SGF.B.createHopToExecutor(loc, args[isolatedIndex].getValue(),
                               /*mandatory*/false);
   }
 
-  // If the input function has caller isolation, we need to fill in that
-  // argument with the formal isolation of the output function.
+  // If the input function has caller isolation (and thus is async), we need to
+  // fill in that argument with the formal isolation of the output function and
+  // hop to it so that it is able to assume that it does not need to hop in its
+  // prologue.
   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: {
-      auto globalActor =
-          outputIsolation.getGlobalActorType()->getCanonicalType();
-      argValues.push_back(
-          SGF.emitGlobalActorIsolation(loc, globalActor).getValue());
-      break;
-    }
-    case FunctionTypeIsolation::Kind::NonIsolatedCaller: {
-      argValues.push_back(implicitIsolationParam.getValue());
-      break;
-    }
-    case FunctionTypeIsolation::Kind::Parameter:
-      // 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");
-    }
+    auto forwardedIsolationValue = [&]() -> SILValue {
+      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.
+        return SGF.emitNonIsolatedIsolation(loc).getValue();
+      case FunctionTypeIsolation::Kind::GlobalActor: {
+        auto globalActor =
+            outputIsolation.getGlobalActorType()->getCanonicalType();
+        return SGF.emitGlobalActorIsolation(loc, globalActor).getValue();
+      }
+      case FunctionTypeIsolation::Kind::NonIsolatedCaller: {
+        return implicitIsolationParam.getValue();
+      }
+      case FunctionTypeIsolation::Kind::Parameter:
+        // 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::report_fatal_error("Should never see this?!");
+      }
+    }();
+
+    assert(forwardedIsolationValue);
+    SGF.B.createHopToExecutor(loc, forwardedIsolationValue, false);
+    argValues.push_back(forwardedIsolationValue);
   }
 
   // Add the rest of the arguments.
@@ -5649,6 +5654,16 @@ static void buildThunkBody(SILGenFunction &SGF, SILLocation loc,
   auto fun = fnType->isCalleeGuaranteed() ? fnValue.borrow(SGF, loc).getValue()
                                           : fnValue.forward(SGF);
 
+  // If our thunk is a caller isolated function, we need to insert a hop to
+  // return to our isolation so that our caller can assume that we preserve
+  // isolation.
+  if (auto isolatedArg = llvm::cast_or_null<SILFunctionArgument>(
+          SGF.F.maybeGetIsolatedArgument());
+      isolatedArg && isolatedArg->getKnownParameterInfo().hasOption(
+                         SILParameterInfo::ImplicitLeading)) {
+    SGF.emitScopedHopToTargetActor(loc, isolatedArg);
+  }
+
   SILValue innerResult =
       SGF.emitApplyWithRethrow(loc, fun,
                                /*substFnType*/ fnValue.getType(),
@@ -7080,6 +7095,15 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
   if (options.contains(ThunkGenFlag::CalleeHasImplicitIsolatedParam)) {
     auto baseIsolation =
         swift::getActorIsolation(base.getAbstractFunctionDecl());
+    std::optional<ActorIsolation> derivedIsolationCache;
+    auto getDerivedIsolation = [&]() -> ActorIsolation {
+      if (!derivedIsolationCache) {
+        derivedIsolationCache =
+            swift::getActorIsolation(derived.getAbstractFunctionDecl());
+      }
+      return *derivedIsolationCache;
+    };
+
     switch (baseIsolation) {
     case ActorIsolation::Unspecified:
     case ActorIsolation::Nonisolated:
@@ -7096,8 +7120,7 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
     }
     case ActorIsolation::ActorInstance:
     case ActorIsolation::CallerIsolationInheriting: {
-      auto derivedIsolation =
-          swift::getActorIsolation(derived.getAbstractFunctionDecl());
+      auto derivedIsolation = getDerivedIsolation();
       switch (derivedIsolation) {
       case ActorIsolation::Unspecified:
       case ActorIsolation::Nonisolated:
@@ -7124,6 +7147,14 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
       break;
     }
     }
+
+    // If our derived isolation is caller isolation inheriting and our base
+    // isn't, we need to insert a hop so that derived can assume that it does
+    // not have to hop in its prologue.
+    if (!baseIsolation.isCallerIsolationInheriting() &&
+        getDerivedIsolation().isCallerIsolationInheriting()) {
+      B.createHopToExecutor(loc, args.back(), false /*mandatory*/);
+    }
   }
 
   // Then, the arguments.
@@ -7142,6 +7173,15 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
     derivedRef = B.createFunctionRefFor(loc, implFn);
   }
 
+  // Check if our base was caller isolation inheriting. In such a case, we need
+  // to insert a hop back to our original actor.
+  if (auto isolatedArg = llvm::dyn_cast_or_null<SILFunctionArgument>(
+          F.maybeGetIsolatedArgument());
+      isolatedArg && isolatedArg->getKnownParameterInfo().hasOption(
+                         SILParameterInfo::ImplicitLeading)) {
+    emitScopedHopToTargetActor(loc, {isolatedArg});
+  }
+
   SILValue result;
 
   switch (coroutineKind) {
@@ -7192,7 +7232,7 @@ SILGenFunction::emitVTableThunk(SILDeclRef base,
     result = B.createTuple(loc, {});
     break;
   }
-  
+
   scope.pop();
   B.createReturn(loc, result);
 
@@ -7560,6 +7600,14 @@ void SILGenFunction::emitProtocolWitness(
   if (options.contains(ThunkGenFlag::CalleeHasImplicitIsolatedParam)) {
     auto reqtIsolation =
         swift::getActorIsolation(requirement.getAbstractFunctionDecl());
+    std::optional<ActorIsolation> witnessIsolationCache;
+    auto getWitnessIsolation = [&]() -> ActorIsolation {
+      if (!witnessIsolationCache) {
+        witnessIsolationCache =
+            swift::getActorIsolation(witness.getAbstractFunctionDecl());
+      }
+      return *witnessIsolationCache;
+    };
     switch (reqtIsolation) {
     case ActorIsolation::Unspecified:
     case ActorIsolation::Nonisolated:
@@ -7576,8 +7624,7 @@ void SILGenFunction::emitProtocolWitness(
     }
     case ActorIsolation::ActorInstance:
     case ActorIsolation::CallerIsolationInheriting: {
-      auto witnessIsolation =
-          swift::getActorIsolation(witness.getAbstractFunctionDecl());
+      auto witnessIsolation = getWitnessIsolation();
       switch (witnessIsolation) {
       case ActorIsolation::Unspecified:
       case ActorIsolation::Nonisolated:
@@ -7604,6 +7651,25 @@ void SILGenFunction::emitProtocolWitness(
       break;
     }
     }
+
+    // If our reqtIsolation was not caller isolation inheriting, but our witness
+    // isolation is caller isolation inheriting, hop onto the reqtIsolation so
+    // that it is safe for our witness to assume that it is already on its
+    // actor.
+    if (!reqtIsolation.isCallerIsolationInheriting() &&
+        getWitnessIsolation().isCallerIsolationInheriting()) {
+      B.createHopToExecutor(loc, args.back(), false /*mandatory*/);
+    }
+  }
+
+  // Check if our caller has an implicit isolated param. If so, we need to
+  // insert a hop at the end of the function to ensure that our caller can
+  // assume that we are going to be on its isolation.
+  if (auto isolatedArg =
+          cast_or_null<SILFunctionArgument>(F.maybeGetIsolatedArgument());
+      isolatedArg && isolatedArg->getKnownParameterInfo().hasOption(
+                         SILParameterInfo::ImplicitLeading)) {
+    emitScopedHopToTargetActor(loc, SILValue(isolatedArg));
   }
 
   //   - the rest of the arguments

test/Concurrency/attr_execution/classes_silgen.swift

@@ -6,8 +6,7 @@
 
 // We should only produce thunks when going to/from nonisolated(nonsending)
 // since that is the only thing that makes a true ABI change since we have an
-// extra parameter. For isolation purposes, we can rely on async functions
-// hopping in their prolog and not need a thunk for the purposes of isolation changing.
+// extra parameter.
 
 // NOTE: We use implicit-check-not to make sure we do not create any other
 // thunks beyond the ones we pattern match.
@@ -31,6 +30,7 @@ class AllConcurrent : SuperKlass {
   // CHECK: bb0([[ACTOR:%.*]] : @guaranteed $Optional<any Actor>, [[PARAM:%.*]] : @guaranteed $AllConcurrent):
   // CHECK:  [[FUNC:%.*]] = function_ref @$s21attr_execution_silgen13AllConcurrentC10callerTestyyYaF : $@convention(method) @async (@guaranteed AllConcurrent) -> ()
   // CHECK:  apply [[FUNC]]([[PARAM]])
+  // CHECK:  hop_to_executor [[ACTOR]]
   // CHECK: } // end sil function '$s21attr_execution_silgen13AllConcurrentC10callerTestyyYaFAA10SuperKlassCADyyYaFTV'
   @concurrent override func callerTest() async {}
   @concurrent override func concurrentTest() async {}
@@ -44,6 +44,7 @@ class AllNonIsolatedUnsafe : SuperKlass {
   // CHECK-NEXT: sil private [thunk] [ossa] @$s21attr_execution_silgen20AllNonIsolatedUnsafeC14concurrentTestyyYaFAA10SuperKlassCADyyYaFTV : $@convention(method) @async (@guaranteed AllNonIsolatedUnsafe) -> () {
   // CHECK: bb0([[ARG:%.*]] : @guaranteed
   // CHECK:   [[ACTOR:%.*]] = enum $Optional<any Actor>, #Optional.none!enumelt
+  // CHECK:   hop_to_executor [[ACTOR]]
   // CHECK:   [[FUNC:%.*]] = function_ref @$s21attr_execution_silgen20AllNonIsolatedUnsafeC14concurrentTestyyYaF : $@convention(method) @async (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>, @guaranteed AllNonIsolatedUnsafe) -> ()
   // CHECK:   apply [[FUNC]]([[ACTOR]], [[ARG]])
   // CHECK: } // end sil function '$s21attr_execution_silgen20AllNonIsolatedUnsafeC14concurrentTestyyYaFAA10SuperKlassCADyyYaFTV'
@@ -55,6 +56,7 @@ class AllNonIsolatedUnsafe : SuperKlass {
   // CHECK:   [[ACTOR:%.*]] = apply {{%.*}}({{%.*}}) : $@convention(method) (@thick MainActor.Type) -> @owned MainActor
   // CHECK:   [[ACTOR_E:%.*]] = init_existential_ref [[ACTOR]]
   // CHECK:   [[ACTOR_E_OPT:%.*]] = enum $Optional<any Actor>, #Optional.some!enumelt, [[ACTOR_E]]
+  // CHECK:   hop_to_executor [[ACTOR_E_OPT]]
   // CHECK:   [[FUNC:%.*]] = function_ref @$s21attr_execution_silgen20AllNonIsolatedUnsafeC13mainActorTestyyYaF : $@convention(method) @async (@sil_isolated @sil_implicit_leading_param @guaranteed Optional<any Actor>, @guaranteed AllNonIsolatedUnsafe) -> ()
   // CHECK:   apply [[FUNC]]([[ACTOR_E_OPT]], [[ARG]])
   // CHECK: } // end sil function '$s21attr_execution_silgen20AllNonIsolatedUnsafeC13mainActorTestyyYaFAA10SuperKlassCADyyYaFTV'