Merge pull request #71099 from gottesmm/pr-387f6b4fba6f40c5cb78dd54f3fd718902f9d38d

[region-isolation] Make it so that we only propagate actor self if the callee takes self as isolated.

Author: gottesmm

include/swift/AST/ASTBridging.h

@@ -1206,7 +1206,8 @@ enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedTypeAttrKind : size_t {
   BridgedTypeAttrKind_moveOnly,
   BridgedTypeAttrKind_thin,
   BridgedTypeAttrKind_thick,
-  BridgedTypeAttrKind_Count
+  BridgedTypeAttrKind_Count,
+  BridgedTypeAttrKind_isolated,
 };
 
 SWIFT_NAME("BridgedTypeAttrKind.init(from:)")

include/swift/AST/Attr.def

@@ -99,6 +99,7 @@ TYPE_ATTR(yield_many)
 TYPE_ATTR(captures_generics)
 // Used at the SIL level to mark a type as moveOnly.
 TYPE_ATTR(moveOnly)
+TYPE_ATTR(isolated)
 
 // SIL metatype attributes.
 TYPE_ATTR(thin)

include/swift/AST/TypeDifferenceVisitor.h

@@ -291,7 +291,7 @@ class CanTypeDifferenceVisitor : public CanTypePairVisitor<Impl, bool> {
   bool visitComponent(CanType type1, CanType type2,
                       SILParameterInfo param1, SILParameterInfo param2) {
     if (param1.getConvention() != param2.getConvention() ||
-        param1.getOptions().containsOnly(param2.getOptions()))
+        !param1.getOptions().containsOnly(param2.getOptions()))
       return asImpl().visitDifferentTypeStructure(type1, type2);
 
     return asImpl().visit(param1.getInterfaceType(),
@@ -410,6 +410,7 @@ class CanTypeDifferenceVisitor : public CanTypePairVisitor<Impl, bool> {
     return asImpl().visit(CanType(componentType1), CanType(componentType2));
   }
 
+protected:
   template <class T>
   bool visitComponentArray(CanType type1, CanType type2, T array1, T array2) {
     if (array1.size() != array2.size())

include/swift/AST/Types.h

@@ -4129,6 +4129,9 @@ class SILParameterInfo {
 
     /// Set if the given parameter is transferring.
     Transferring = 0x2,
+
+    /// Set if the given parameter is isolated.
+    Isolated = 0x4,
   };
 
   using Options = OptionSet<Flag>;

lib/AST/ASTPrinter.cpp

@@ -8086,6 +8086,11 @@ void SILParameterInfo::print(ASTPrinter &Printer,
     Printer << "@transferring ";
   }
 
+  if (options.contains(SILParameterInfo::Isolated)) {
+    options -= SILParameterInfo::Isolated;
+    Printer << "@isolated ";
+  }
+
   // If we did not handle a case in Options, this code was not updated
   // appropriately.
   assert(!bool(options) && "Code not updated for introduced option");

lib/AST/TypeRepr.cpp

@@ -236,6 +236,8 @@ void AttributedTypeRepr::printAttrs(ASTPrinter &Printer,
     Printer.printSimpleAttr("@Sendable") << " ";
   if (hasAttr(TAK_noDerivative))
     Printer.printSimpleAttr("@noDerivative") << " ";
+  if (hasAttr(TAK_isolated))
+    Printer.printSimpleAttr("@isolated") << " ";
 
   if (hasAttr(TAK_differentiable)) {
     Printer.callPrintStructurePre(PrintStructureKind::BuiltinAttribute);

lib/ASTGen/Sources/ASTGen/Types.swift

@@ -415,7 +415,8 @@ extension ASTGenVisitor {
           .sil_weak, .sil_unowned, .inout, .block_storage, .box,
           .dynamic_self, .sil_unmanaged, .error, .error_indirect,
           .error_unowned, .direct, .inout_aliasable,
-          .in_guaranteed, .in_constant, .captures_generics, .moveOnly:
+          .in_guaranteed, .in_constant, .captures_generics, .moveOnly,
+          .isolated:
           fallthrough
 
         case .autoclosure, .escaping, .noescape, .noDerivative, .async,

lib/SIL/IR/SILFunctionType.cpp

@@ -1737,6 +1737,8 @@ class DestructureInputs {
       param = param.addingOption(SILParameterInfo::NotDifferentiable);
     if (origFlags.isTransferring())
       param = param.addingOption(SILParameterInfo::Transferring);
+    if (origFlags.isIsolated())
+      param = param.addingOption(SILParameterInfo::Isolated);
 
     Inputs.push_back(param);
     maybeAddForeignParameters();

lib/SILGen/SILGenThunk.cpp

@@ -32,6 +32,7 @@
 #include "swift/AST/ForeignAsyncConvention.h"
 #include "swift/AST/ForeignErrorConvention.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/TypeDifferenceVisitor.h"
 #include "swift/Basic/STLExtras.h"
 #include "swift/SIL/FormalLinkage.h"
 #include "swift/SIL/PrettyStackTrace.h"
@@ -139,6 +140,51 @@ void SILGenModule::emitBackDeploymentThunk(SILDeclRef thunk) {
   emitFunctionDefinition(thunk, getFunction(thunk, ForDefinition));
 }
 
+namespace {
+
+/// Checker that validates that a distributed thunk is completely the same
+/// except that self can vary by isolation.
+struct DistributedThunkDiffChecker
+    : CanTypeDifferenceVisitor<DistributedThunkDiffChecker> {
+  using SuperTy = CanTypeDifferenceVisitor<DistributedThunkDiffChecker>;
+
+  bool visitSILFunctionTypeComponents(CanSILFunctionType type1,
+                                      CanSILFunctionType type2) {
+    // If they do not both have a self param. Just delegate to our parent.
+    if (!type1->hasSelfParam() || !type2->hasSelfParam()) {
+      return SuperTy::visitSILFunctionTypeComponents(type1, type2);
+    }
+
+    // Otherwise, we both have self. First check if we have the same number of
+    // parameters.
+    auto type1Params = type1->getParameters();
+    auto type2Params = type2->getParameters();
+    if (type1Params.size() != type2Params.size())
+      return visitDifferentTypeStructure(type1, type2);
+
+    // Then check if self is the same ignoring isolation.
+    auto self1 = type1Params.back();
+    auto self2 = type2Params.back();
+    auto self1Options = self1.getOptions() - SILParameterInfo::Isolated;
+    auto self2Options = self2.getOptions() - SILParameterInfo::Isolated;
+
+    if (self1.getConvention() != self2.getConvention() ||
+        !self1Options.containsOnly(self2Options))
+      return visitDifferentTypeStructure(type1, type2);
+
+    // Finally, check our self type, non-self components, results, and yields.
+    return visit(self1.getInterfaceType(), self2.getInterfaceType()) ||
+           visitComponentArray(type1, type2, type1Params.drop_back(),
+                               type2Params.drop_back()) ||
+           visitComponentArray(type1, type2, type1->getResults(),
+                               type2->getResults()) ||
+           visitComponentArray(type1, type2, type1->getYields(),
+                               type2->getYields());
+  }
+};
+
+} // namespace
+
 SILValue
 SILGenFunction::emitGlobalFunctionRef(SILLocation loc, SILDeclRef constant,
                                       SILConstantInfo constantInfo,
@@ -171,18 +217,35 @@ SILGenFunction::emitGlobalFunctionRef(SILLocation loc, SILDeclRef constant,
   auto existingType =
       f->getLoweredFunctionTypeInContext(B.getTypeExpansionContext());
   if (existingType != constantFnTypeInContext) {
-    // This can happen for example when using @_silgen_name or @_extern(c)
-    // attributes
-    SGM.diagnose(loc.getSourceLoc(), diag::function_type_mismatch, existingType,
-                 constantFnTypeInContext);
-    SGM.diagnose(f->getLocation().getSourceLoc(), diag::function_declared_here);
-    return SILUndef::get(constantInfo.getSILType(), F);
+    auto emitError = [&] {
+      // This can happen for example when using @_silgen_name or @_extern(c)
+      // attributes
+      SGM.diagnose(loc.getSourceLoc(), diag::function_type_mismatch,
+                   existingType, constantFnTypeInContext);
+      SGM.diagnose(f->getLocation().getSourceLoc(),
+                   diag::function_declared_here);
+      return SILUndef::get(constantInfo.getSILType(), F);
+    };
+
+    // If we have a distributed thunk, see if we only differ by isolation.
+    if (f->isDistributed() && f->isThunk()) {
+      DistributedThunkDiffChecker diffChecker;
+      if (diffChecker.visit(existingType, constantFnTypeInContext)) {
+        return emitError();
+      }
+
+      // We differ only by isolation... so do not error.
+    } else {
+      // This can happen for example when using @_silgen_name or @_extern(c)
+      // attributes
+      return emitError();
+    }
   }
 
   if (callPreviousDynamicReplaceableImpl)
     return B.createPreviousDynamicFunctionRef(loc, f);
-  else
-    return B.createFunctionRefFor(loc, f);
+
+  return B.createFunctionRefFor(loc, f);
 }
 
 static const clang::Type *prependParameterType(

lib/SILOptimizer/Analysis/RegionAnalysis.cpp

@@ -1619,10 +1619,14 @@ class PartitionOpTranslator {
 
   void translateNonIsolationCrossingSILApply(FullApplySite fas) {
     SILMultiAssignOptions options;
+
+    // If self is an actor and we are isolated to it, propagate actor self.
     if (fas.hasSelfArgument()) {
-      if (auto self = fas.getSelfArgument()) {
-        if (self->getType().isActor())
-          options |= SILMultiAssignFlags::PropagatesActorSelf;
+      auto &self = fas.getSelfArgumentOperand();
+      if (self.get()->getType().isActor() &&
+          fas.getArgumentParameterInfo(self).hasOption(
+              SILParameterInfo::Isolated)) {
+        options |= SILMultiAssignFlags::PropagatesActorSelf;
       }
     }