[region-isolation] A few improvements + If a value is dynamically actor isolated, do not consider it transferred if the transfer statically was to that same actor isolation.

include/swift/AST/DiagnosticsSIL.def

@@ -987,7 +987,7 @@ NOTE(regionbasedisolation_named_stronglytransferred_binding, none,
       "%0 used after being passed as a transferring parameter; Later uses could race",
       (Identifier))
 NOTE(regionbasedisolation_named_isolated_closure_yields_race, none,
-      "%0 %1 is captured by %2 closure. %2 uses in closure may race against later %3 uses",
+      "%0 %1 is captured by a %2 closure. %2 uses in closure may race against later %3 uses",
       (StringRef, Identifier, ActorIsolation, ActorIsolation))
 
 // Misc Error.

include/swift/SIL/SILFunction.h

@@ -347,6 +347,9 @@ class SILFunction
   /// block indices.
   unsigned BlockListChangeIdx = 0;
 
+  /// The isolation of this function.
+  std::optional<ActorIsolation> actorIsolation;
+
   /// The function's bare attribute. Bare means that the function is SIL-only
   /// and does not require debug info.
   unsigned Bare : 1;
@@ -1367,6 +1370,14 @@ class SILFunction
     return false;
   }
 
+  void setActorIsolation(ActorIsolation newActorIsolation) {
+    actorIsolation = newActorIsolation;
+  }
+
+  std::optional<ActorIsolation> getActorIsolation() const {
+    return actorIsolation;
+  }
+
   //===--------------------------------------------------------------------===//
   // Block List Access
   //===--------------------------------------------------------------------===//

include/swift/SILOptimizer/Analysis/RegionAnalysis.h

@@ -286,6 +286,13 @@ class regionanalysisimpl::TrackableValue {
   bool isTransferringParameter() const;
 
   void print(llvm::raw_ostream &os) const {
+    os << "TrackableValue. State: ";
+    valueState.print(os);
+    os << "\n    Rep Value: ";
+    getRepresentative().print(os);
+  }
+
+  void printVerbose(llvm::raw_ostream &os) const {
     os << "TrackableValue. State: ";
     valueState.print(os);
     os << "\n    Rep Value: " << getRepresentative();

include/swift/SILOptimizer/Utils/PartitionUtils.h

@@ -17,6 +17,7 @@
 #include "swift/Basic/FrozenMultiMap.h"
 #include "swift/Basic/ImmutablePointerSet.h"
 #include "swift/Basic/LLVM.h"
+#include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILInstruction.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
@@ -204,7 +205,7 @@ class SILIsolationInfo {
       return getActorIsolated(actorIsolation);
     if (nomDecl->isActor())
       return {Kind::Actor, nomDecl};
-    return {};
+    return SILIsolationInfo();
   }
 
   static SILIsolationInfo getTaskIsolated(SILValue value) {
@@ -756,8 +757,8 @@ struct PartitionOpEvaluator {
       // value... emit an error.
       if (auto *transferredOperandSet = p.getTransferred(op.getOpArgs()[1])) {
         for (auto transferredOperand : transferredOperandSet->data()) {
-          handleLocalUseAfterTransfer(op, op.getOpArgs()[1],
-                                      transferredOperand);
+          handleLocalUseAfterTransferHelper(op, op.getOpArgs()[1],
+                                            transferredOperand);
         }
       }
       p.assignElement(op.getOpArgs()[0], op.getOpArgs()[1]);
@@ -789,6 +790,21 @@ struct PartitionOpEvaluator {
       std::tie(transferredRegionIsolation, isClosureCapturedElt) =
           getIsolationRegionInfo(transferredRegion, op.getSourceOp());
 
+      // Before we do anything, see if our dynamic isolation kind is the same as
+      // the isolation info for our partition op. If they match, this is not a
+      // real transfer operation.
+      //
+      // DISCUSSION: We couldn't not emit this earlier since we needed the
+      // dynamic isolation info of our value.
+      if (transferredRegionIsolation.isActorIsolated()) {
+        if (auto calleeIsolationInfo =
+                SILIsolationInfo::get(op.getSourceInst())) {
+          if (transferredRegionIsolation == calleeIsolationInfo) {
+            return;
+          }
+        }
+      }
+
       // If we merged anything, we need to handle a transfer
       // non-transferrable. We pass in the dynamic isolation region info of our
       // region.
@@ -824,14 +840,14 @@ struct PartitionOpEvaluator {
       // if attempting to merge a transferred region, handle the failure
       if (auto *transferredOperandSet = p.getTransferred(op.getOpArgs()[0])) {
         for (auto transferredOperand : transferredOperandSet->data()) {
-          handleLocalUseAfterTransfer(op, op.getOpArgs()[0],
-                                      transferredOperand);
+          handleLocalUseAfterTransferHelper(op, op.getOpArgs()[0],
+                                            transferredOperand);
         }
       }
       if (auto *transferredOperandSet = p.getTransferred(op.getOpArgs()[1])) {
         for (auto transferredOperand : transferredOperandSet->data()) {
-          handleLocalUseAfterTransfer(op, op.getOpArgs()[1],
-                                      transferredOperand);
+          handleLocalUseAfterTransferHelper(op, op.getOpArgs()[1],
+                                            transferredOperand);
         }
       }
 
@@ -844,8 +860,8 @@ struct PartitionOpEvaluator {
              "Require PartitionOp's argument should already be tracked");
       if (auto *transferredOperandSet = p.getTransferred(op.getOpArgs()[0])) {
         for (auto transferredOperand : transferredOperandSet->data()) {
-          handleLocalUseAfterTransfer(op, op.getOpArgs()[0],
-                                      transferredOperand);
+          handleLocalUseAfterTransferHelper(op, op.getOpArgs()[0],
+                                            transferredOperand);
         }
       }
       return;
@@ -858,6 +874,43 @@ struct PartitionOpEvaluator {
     for (auto &o : ops)
       apply(o);
   }
+
+  /// Provides a way for subclasses to disable the error squelching
+  /// functionality.
+  ///
+  /// Used by the unittests.
+  bool shouldTryToSquelchErrors() const {
+    return asImpl().shouldTryToSquelchErrors();
+  }
+
+private:
+  // Private helper that squelches the error if our transfer instruction and our
+  // use have the same isolation.
+  void
+  handleLocalUseAfterTransferHelper(const PartitionOp &op, Element elt,
+                                    TransferringOperand *transferringOp) const {
+    if (shouldTryToSquelchErrors()) {
+      if (auto isolationInfo = SILIsolationInfo::get(op.getSourceInst())) {
+        if (isolationInfo.isActorIsolated() &&
+            isolationInfo == SILIsolationInfo::get(transferringOp->getUser()))
+          return;
+      }
+
+      // If our instruction does not have any isolation info associated with it,
+      // it must be nonisolated. See if our function has a matching isolation to
+      // our transferring operand. If so, we can squelch this.
+      if (auto functionIsolation =
+              transferringOp->getUser()->getFunction()->getActorIsolation()) {
+        if (functionIsolation->isActorIsolated() &&
+            SILIsolationInfo::getActorIsolated(*functionIsolation) ==
+                SILIsolationInfo::get(transferringOp->getUser()))
+          return;
+      }
+    }
+
+    // Ok, we actually need to emit a call to the callback.
+    return handleLocalUseAfterTransfer(op, elt, transferringOp);
+  }
 };
 
 /// A base implementation that can be used to default initialize CRTP
@@ -927,6 +980,9 @@ struct PartitionOpEvaluatorBaseImpl : PartitionOpEvaluator<Subclass> {
   /// to access the instruction in the evaluator which creates a problem when
   /// since the operand we pass in is a dummy operand.
   bool isClosureCaptured(Element elt, Operand *op) const { return false; }
+
+  /// By default squelch errors.
+  bool shouldTryToSquelchErrors() const { return true; }
 };
 
 /// A subclass of PartitionOpEvaluatorBaseImpl that doesn't have any special

lib/SILGen/SILGen.cpp

@@ -1217,6 +1217,10 @@ void SILGenModule::preEmitFunction(SILDeclRef constant, SILFunction *F,
   if (F->getLoweredFunctionType()->isPolymorphic())
     F->setGenericEnvironment(Types.getConstantGenericEnvironment(constant));
 
+  // Set the actor isolation of the function to its innermost decl context.
+  F->setActorIsolation(
+      getActorIsolationOfContext(constant.getInnermostDeclContext()));
+
   // Create a debug scope for the function using astNode as source location.
   F->setDebugScope(new (M) SILDebugScope(Loc, F));
 

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1470,6 +1470,7 @@ class PartitionOpTranslator {
     llvm::SmallVector<Element, 8> nonSendableJoinedIndices;
     llvm::SmallVector<Element, 8> nonSendableSeparateIndices;
     for (SILArgument *arg : functionArguments) {
+      // This will decide what the isolation region is.
       if (auto state = tryToTrackValue(arg)) {
         // If we can transfer our parameter, just add it to
         // nonSendableSeparateIndices.
@@ -1485,9 +1486,8 @@ class PartitionOpTranslator {
 
         // Otherwise, it is one of our merged parameters. Add it to the never
         // transfer list and to the region join list.
-        LLVM_DEBUG(llvm::dbgs() << "    %%" << state->getID() << ": " << *arg);
-        valueMap.mergeIsolationRegionInfo(
-            arg, SILIsolationInfo::getTaskIsolated(arg));
+        LLVM_DEBUG(llvm::dbgs() << "    %%" << state->getID() << ": ";
+                   state->print(llvm::dbgs()); llvm::dbgs() << *arg);
         nonSendableJoinedIndices.push_back(state->getID());
       }
     }
@@ -3460,7 +3460,9 @@ void RegionAnalysisValueMap::print(llvm::raw_ostream &os) const {
   }
   std::sort(temp.begin(), temp.end());
   for (auto p : temp) {
-    os << "%%" << p.first << ": " << p.second;
+    os << "%%" << p.first << ": ";
+    auto value = getValueForId(Element(p.first));
+    value->print(os);
   }
 #endif
 }

lib/SILOptimizer/Mandatory/TransferNonSendable.cpp

@@ -1021,11 +1021,19 @@ class TransferNonTransferrableDiagnosticEmitter {
         .highlight(getOperand()->getUser()->getLoc().getSourceRange());
   }
 
-  void emitFunctionArgumentClosure(SourceLoc loc, Type type,
-                                   ApplyIsolationCrossing crossing) {
-    diagnoseError(loc, diag::regionbasedisolation_arg_transferred,
-                  "task-isolated", type, crossing.getCalleeIsolation())
-        .highlight(getOperand()->getUser()->getLoc().getSourceRange());
+  void emitNamedFunctionArgumentClosure(SILLocation loc, Identifier name,
+                                        ApplyIsolationCrossing crossing) {
+    emitNamedOnlyError(loc, name);
+    SmallString<64> descriptiveKindStr;
+    {
+      llvm::raw_svector_ostream os(descriptiveKindStr);
+      getIsolationRegionInfo().printForDiagnostics(os);
+    }
+    diagnoseNote(loc,
+                 diag::regionbasedisolation_named_isolated_closure_yields_race,
+                 descriptiveKindStr, name, crossing.getCalleeIsolation(),
+                 crossing.getCallerIsolation())
+        .highlight(loc.getSourceRange());
   }
 
   void emitFunctionArgumentApplyStronglyTransferred(SILLocation loc,
@@ -1160,9 +1168,9 @@ bool TransferNonTransferrableDiagnosticInferrer::initForIsolatedPartialApply(
   unsigned opIndex = ApplySite(op->getUser()).getAppliedArgIndex(*op);
   for (auto &p : foundCapturedIsolationCrossing) {
     if (std::get<1>(p) == opIndex) {
-      Type type = std::get<0>(p).getDecl()->getInterfaceType();
-      diagnosticEmitter.emitFunctionArgumentClosure(std::get<0>(p).getLoc(),
-                                                    type, std::get<2>(p));
+      auto loc = RegularLocation(std::get<0>(p).getLoc());
+      diagnosticEmitter.emitNamedFunctionArgumentClosure(
+          loc, std::get<0>(p).getDecl()->getBaseIdentifier(), std::get<2>(p));
       return true;
     }
   }

lib/SILOptimizer/Utils/PartitionUtils.cpp

@@ -42,13 +42,21 @@ static llvm::cl::opt<bool, true> // The parser
 //===----------------------------------------------------------------------===//
 
 SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
-  if (ApplyExpr *apply = inst->getLoc().getAsASTNode<ApplyExpr>())
-    if (auto crossing = apply->getIsolationCrossing())
-      return SILIsolationInfo::getActorIsolated(crossing->getCalleeIsolation());
+  if (ApplyExpr *apply = inst->getLoc().getAsASTNode<ApplyExpr>()) {
+    if (auto crossing = apply->getIsolationCrossing()) {
+      if (crossing->getCalleeIsolation().isActorIsolated())
+        return SILIsolationInfo::getActorIsolated(
+            crossing->getCalleeIsolation());
+    }
+  }
 
   if (auto fas = FullApplySite::isa(inst)) {
-    if (auto crossing = fas.getIsolationCrossing())
-      return SILIsolationInfo::getActorIsolated(crossing->getCalleeIsolation());
+    if (auto crossing = fas.getIsolationCrossing()) {
+      if (crossing->getCalleeIsolation().isActorIsolated()) {
+        return SILIsolationInfo::getActorIsolated(
+            crossing->getCalleeIsolation());
+      }
+    }
 
     if (fas.hasSelfArgument()) {
       auto &self = fas.getSelfArgumentOperand();
@@ -79,9 +87,18 @@ SILIsolationInfo SILIsolationInfo::get(SILInstruction *inst) {
 }
 
 SILIsolationInfo SILIsolationInfo::get(SILFunctionArgument *arg) {
+  // If we have self and our function is actor isolated, all of our arguments
+  // should be marked as actor isolated.
   if (auto *self = arg->getFunction()->maybeGetSelfArgument()) {
-    if (auto *nomDecl = self->getType().getNominalOrBoundGenericNominal()) {
-      return SILIsolationInfo::getActorIsolated(nomDecl);
+    if (auto functionIsolation = arg->getFunction()->getActorIsolation()) {
+      if (functionIsolation->isActorIsolated()) {
+        if (auto *nomDecl = self->getType().getNominalOrBoundGenericNominal()) {
+          if (auto isolationInfo =
+                  SILIsolationInfo::getActorIsolated(nomDecl)) {
+            return isolationInfo;
+          }
+        }
+      }
     }
   }
 
@@ -171,7 +188,7 @@ bool SILIsolationInfo::operator==(const SILIsolationInfo &other) const {
 
     // Otherwise, try to use the inferred actor decl.
     auto *lhsDecl = tryInferActorDecl();
-    auto *rhsDecl = tryInferActorDecl();
+    auto *rhsDecl = other.tryInferActorDecl();
     if (lhsDecl && rhsDecl)
       return lhsDecl == rhsDecl;