Merge pull request #83177 from gottesmm/rdar155905383-finished

[concurrency] Add fixes from rjmccall's review on 6.2 to main

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

lib/SILOptimizer/Mandatory/OptimizeHopToExecutor.cpp

@@ -342,11 +342,12 @@ void OptimizeHopToExecutor::updateNeedExecutor(int &needExecutor,
   // 6.2, since caller isolation inheriting functions will no longer be
   // considered suspension points, we will be able to sink this code into needs
   // executor.
-  if (auto isolation = inst->getFunction()->getActorIsolation();
-      isolation && isolation->isCallerIsolationInheriting() &&
-      isa<ReturnInst>(inst)) {
-    needExecutor = BlockState::ExecutorNeeded;
-    return;
+  if (isa<ReturnInst>(inst)) {
+    if (auto isolation = inst->getFunction()->getActorIsolation();
+        isolation && isolation->isCallerIsolationInheriting()) {
+      needExecutor = BlockState::ExecutorNeeded;
+      return;
+    }
   }
 
   if (isSuspensionPoint(inst)) {

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'