[Distributed] Implement whenLocal properly

include/swift/AST/Attr.def

@@ -728,7 +728,7 @@ DECL_ATTR(_backDeploy, BackDeploy,
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIBreakingToRemove,
   129)
 
-CONTEXTUAL_SIMPLE_DECL_ATTR(_local, KnownToBeLocal,
+CONTEXTUAL_SIMPLE_DECL_ATTR(_local, DistributedKnownToBeLocal,
   DeclModifier | OnFunc | OnParam | OnVar |
   UserInaccessible |
   ABIBreakingToAdd | ABIBreakingToRemove |

include/swift/AST/Decl.h

@@ -4829,6 +4829,9 @@ class AbstractStorageDecl : public ValueDecl {
   }
   bool isCompileTimeConst() const;
 
+  /// Is a reference to a distributed actor that is "known to be local".
+  bool isDistributedKnownToBeLocal() const;
+
   /// \returns the way 'static'/'class' should be spelled for this declaration.
   StaticSpellingKind getCorrectStaticSpelling() const;
 
@@ -5359,10 +5362,6 @@ class VarDecl : public AbstractStorageDecl {
   /// 'distributed' and and nullptr otherwise.
   FuncDecl *getDistributedThunk() const;
 
-  /// Is this var known to be a "local" distributed actor,
-  /// if so the implicit throwing ans some isolation checks can be skipped.
-  bool isKnownToBeLocal() const;
-
   /// Is this a stored property that will _not_ trigger any user-defined code
   /// upon any kind of access?
   bool isOrdinaryStoredProperty() const;
@@ -5635,9 +5634,13 @@ class ParamDecl : public VarDecl {
 
     /// Whether or not this parameter is '_const'.
     IsCompileTimeConst = 1 << 1,
+
+    // FIXME: it's not on the name but on the type
+    /// Whether or not this parameter is a '_local DistributedActor' reference.
+    IsDistributedKnownToBeLocal = 1 << 2,
   };
 
-  llvm::PointerIntPair<Identifier, 2, OptionSet<ArgumentNameFlags>>
+  llvm::PointerIntPair<Identifier, 3, OptionSet<ArgumentNameFlags>>
       ArgumentNameAndFlags;
   SourceLoc ParameterNameLoc;
   SourceLoc ArgumentNameLoc;
@@ -5878,6 +5881,21 @@ class ParamDecl : public VarDecl {
                                       : flags - Flags::IsIsolated);
   }
 
+  /// Whether or not this parameter is marked with 'isolated'.
+  bool isDistributedKnownToBeLocal() const {
+    return ArgumentNameAndFlags.getInt().contains(ArgumentNameFlags::IsDistributedKnownToBeLocal);
+  }
+
+  void setDistributedKnownToBeLocal(bool value = true) {
+    auto flags = ArgumentNameAndFlags.getInt();
+    if (value && !getAttrs().hasAttribute<DistributedKnownToBeLocalAttr>()) {
+      getAttrs().add(new (getASTContext())
+                         DistributedKnownToBeLocalAttr(/*isImplicit=*/true));
+    }
+    ArgumentNameAndFlags.setInt(value ? flags | ArgumentNameFlags::IsDistributedKnownToBeLocal
+                                      : flags - ArgumentNameFlags::IsDistributedKnownToBeLocal);
+  }
+
   /// Whether or not this parameter is marked with '_const'.
   bool isCompileTimeConst() const {
     return ArgumentNameAndFlags.getInt().contains(

include/swift/AST/DiagnosticsSema.def

@@ -4605,7 +4605,8 @@ ERROR(distributed_actor_isolated_method,none,
       "only 'distributed' instance methods can be called on a potentially remote distributed actor",
       ())
 ERROR(distributed_local_cannot_be_used,none,
-      "'local' cannot be used in user-defined code currently",
+      "'_local' cannot be used in this position, "
+      "only usable by a DistributedActor's 'whenLocal' method",
       ())
 ERROR(distributed_thunk_cannot_be_used,none,
       "'distributed_thunk' cannot be used in user-defined code",
@@ -4667,6 +4668,10 @@ NOTE(protocol_isolated_to_global_actor_here,none,
 
 ERROR(isolated_parameter_not_actor,none,
       "'isolated' parameter has non-actor type %0", (Type))
+ERROR(local_parameter_not_distributed_actor,none,
+      "'_local' parameter has non-distributed-actor type %0", (Type))
+ERROR(local_not_supported_in_pattern,none,
+      "'_local' is not support in patterns", ())
 
 WARNING(non_sendable_param_type,none,
         "non-sendable type %0 %select{passed in call to %4 %2 %3|"

include/swift/AST/Expr.h

@@ -499,6 +499,8 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
   bool printConstExprValue(llvm::raw_ostream *OS, llvm::function_ref<bool(Expr*)> additionalCheck) const;
   bool isSemanticallyConstExpr(llvm::function_ref<bool(Expr*)> additionalCheck = nullptr) const;
 
+  bool isDistributedKnownToBeLocal() const;
+
   /// Returns false if this expression needs to be wrapped in parens when
   /// used inside of a any postfix expression, true otherwise.
   ///

include/swift/AST/TypeRepr.h

@@ -962,7 +962,8 @@ class SpecifierTypeRepr : public TypeRepr {
     return T->getKind() == TypeReprKind::InOut ||
            T->getKind() == TypeReprKind::Shared ||
            T->getKind() == TypeReprKind::Owned ||
-           T->getKind() == TypeReprKind::Isolated;
+           T->getKind() == TypeReprKind::Isolated ||
+           T->getKind() == TypeReprKind::DistributedKnownToBeLocal;
   }
   static bool classof(const SpecifierTypeRepr *T) { return true; }
 
@@ -1048,6 +1049,21 @@ class CompileTimeConstTypeRepr : public SpecifierTypeRepr {
   static bool classof(const CompileTimeConstTypeRepr *T) { return true; }
 };
 
+/// A '_local' type.
+/// \code
+///   x : _local SomeDistributedActor
+/// \endcode
+class DistributedKnownToBeLocalTypeRepr : public SpecifierTypeRepr {
+public:
+  DistributedKnownToBeLocalTypeRepr(TypeRepr *Base, SourceLoc Loc)
+      : SpecifierTypeRepr(TypeReprKind::DistributedKnownToBeLocal, Base, Loc) {}
+
+  static bool classof(const TypeRepr *T) {
+    return T->getKind() == TypeReprKind::DistributedKnownToBeLocal;
+  }
+  static bool classof(const DistributedKnownToBeLocalTypeRepr *T) { return true; }
+};
+
 /// A TypeRepr for a known, fixed type.
 ///
 /// Fixed type representations should be used sparingly, in places
@@ -1348,6 +1364,7 @@ inline bool TypeRepr::isSimple() const {
   case TypeReprKind::Isolated:
   case TypeReprKind::Placeholder:
   case TypeReprKind::CompileTimeConst:
+  case TypeReprKind::DistributedKnownToBeLocal:
     return true;
   }
   llvm_unreachable("bad TypeRepr kind");

include/swift/AST/TypeReprNodes.def

@@ -70,6 +70,7 @@ ABSTRACT_TYPEREPR(Specifier, TypeRepr)
   SPECIFIER_TYPEREPR(Owned, SpecifierTypeRepr)
   SPECIFIER_TYPEREPR(Isolated, SpecifierTypeRepr)
   SPECIFIER_TYPEREPR(CompileTimeConst, SpecifierTypeRepr)
+  SPECIFIER_TYPEREPR(DistributedKnownToBeLocal, SpecifierTypeRepr)
 TYPEREPR(Fixed, TypeRepr)
 TYPEREPR(SILBox, TypeRepr)
 LAST_TYPEREPR(SILBox)

include/swift/AST/Types.h

@@ -2008,7 +2008,8 @@ class ParameterTypeFlags {
     NoDerivative = 1 << 6,
     Isolated     = 1 << 7,
     CompileTimeConst = 1 << 8,
-    NumBits = 9
+    DistributedKnownToBeLocal = 1 << 9,
+    NumBits = 10
   };
   OptionSet<ParameterFlags> value;
   static_assert(NumBits <= 8*sizeof(OptionSet<ParameterFlags>), "overflowed");
@@ -2023,20 +2024,21 @@ class ParameterTypeFlags {
 
   ParameterTypeFlags(bool variadic, bool autoclosure, bool nonEphemeral,
                      ValueOwnership ownership, bool isolated, bool noDerivative,
-                     bool compileTimeConst)
+                     bool compileTimeConst, bool distributedKnownLocal)
       : value((variadic ? Variadic : 0) | (autoclosure ? AutoClosure : 0) |
               (nonEphemeral ? NonEphemeral : 0) |
               uint8_t(ownership) << OwnershipShift |
               (isolated ? Isolated : 0) |
               (noDerivative ? NoDerivative : 0) |
-              (compileTimeConst ? CompileTimeConst : 0)){}
+              (compileTimeConst ? CompileTimeConst : 0) |
+              (distributedKnownLocal ? DistributedKnownToBeLocal : 0)) {}
 
   /// Create one from what's present in the parameter type
   inline static ParameterTypeFlags
   fromParameterType(Type paramTy, bool isVariadic, bool isAutoClosure,
                     bool isNonEphemeral, ValueOwnership ownership,
                     bool isolated, bool isNoDerivative,
-                    bool compileTimeConst);
+                    bool compileTimeConst, bool distributedKnownLocal);
 
   bool isNone() const { return !value; }
   bool isVariadic() const { return value.contains(Variadic); }
@@ -2047,6 +2049,7 @@ class ParameterTypeFlags {
   bool isOwned() const { return getValueOwnership() == ValueOwnership::Owned; }
   bool isIsolated() const { return value.contains(Isolated); }
   bool isCompileTimeConst() const { return value.contains(CompileTimeConst); }
+  bool isDistributedKnownToBeLocal() const { return value.contains(DistributedKnownToBeLocal); }
   bool isNoDerivative() const { return value.contains(NoDerivative); }
 
   ValueOwnership getValueOwnership() const {
@@ -2067,6 +2070,11 @@ class ParameterTypeFlags {
     return ParameterTypeFlags(isConst ? value | ParameterTypeFlags::CompileTimeConst
                                       : value - ParameterTypeFlags::CompileTimeConst);
   }
+
+  ParameterTypeFlags withDistributedKnownToBeLocal(bool isLocal) const {
+    return ParameterTypeFlags(isLocal ? value | ParameterTypeFlags::DistributedKnownToBeLocal
+                                      : value - ParameterTypeFlags::DistributedKnownToBeLocal);
+  }
 
   ParameterTypeFlags withShared(bool isShared) const {
     return withValueOwnership(isShared ? ValueOwnership::Shared
@@ -2175,7 +2183,8 @@ class YieldTypeFlags {
                               /*autoclosure*/ false,
                               /*nonEphemeral*/ false, getValueOwnership(),
                               /*isolated*/ false, /*noDerivative*/ false,
-                              /*compileTimeConst*/false);
+                              /*compileTimeConst*/false,
+                              /*distributedKnownLocal*/false);
   }
 
   bool operator ==(const YieldTypeFlags &other) const {
@@ -2977,6 +2986,9 @@ class AnyFunctionType : public TypeBase {
     /// Whether the parameter is 'isolated'.
     bool isIsolated() const { return Flags.isIsolated(); }
 
+    /// Whether the parameter is '_local'.
+    bool isDistributedKnownToBeLocal() const { return Flags.isDistributedKnownToBeLocal(); }
+
     /// Whether the parameter is 'isCompileTimeConst'.
     bool isCompileTimeConst() const { return Flags.isCompileTimeConst(); }
 
@@ -6755,7 +6767,7 @@ inline TupleTypeElt TupleTypeElt::getWithType(Type T) const {
 inline ParameterTypeFlags ParameterTypeFlags::fromParameterType(
     Type paramTy, bool isVariadic, bool isAutoClosure, bool isNonEphemeral,
     ValueOwnership ownership, bool isolated, bool isNoDerivative,
-    bool compileTimeConst) {
+    bool compileTimeConst, bool distributedKnownLocal) {
   // FIXME(Remove InOut): The last caller that needs this is argument
   // decomposition.  Start by enabling the assertion there and fixing up those
   // callers, then remove this, then remove
@@ -6766,7 +6778,7 @@ inline ParameterTypeFlags ParameterTypeFlags::fromParameterType(
     ownership = ValueOwnership::InOut;
   }
   return {isVariadic, isAutoClosure, isNonEphemeral, ownership, isolated,
-          isNoDerivative, compileTimeConst};
+          isNoDerivative, compileTimeConst, distributedKnownLocal};
 }
 
 inline const Type *BoundGenericType::getTrailingObjectsPointer() const {

include/swift/Demangling/DemangleNodes.def

@@ -338,6 +338,7 @@ NODE(BackDeploymentThunk)
 NODE(BackDeploymentFallback)
 NODE(ExtendedExistentialTypeShape)
 NODE(Uniquable)
+NODE(DistributedKnownToBeLocal)
 
 #undef CONTEXT_NODE
 #undef NODE

include/swift/Parse/Parser.h

@@ -1135,22 +1135,26 @@ class Parser {
                                       SourceLoc &SpecifierLoc,
                                       SourceLoc &IsolatedLoc,
                                       SourceLoc &ConstLoc,
+                                      SourceLoc &DistributedLocalLoc,
                                       TypeAttributes &Attributes) {
     if (Tok.isAny(tok::at_sign, tok::kw_inout) ||
         (Tok.is(tok::identifier) &&
          (Tok.getRawText().equals("__shared") ||
           Tok.getRawText().equals("__owned") ||
           Tok.isContextualKeyword("isolated") ||
-          Tok.isContextualKeyword("_const"))))
+          Tok.isContextualKeyword("_const") ||
+          Tok.isContextualKeyword("_local"))))
       return parseTypeAttributeListPresent(
-          Specifier, SpecifierLoc, IsolatedLoc, ConstLoc, Attributes);
+          Specifier, SpecifierLoc, IsolatedLoc, ConstLoc, DistributedLocalLoc,
+          Attributes);
     return makeParserSuccess();
   }
 
   ParserStatus parseTypeAttributeListPresent(ParamDecl::Specifier &Specifier,
                                              SourceLoc &SpecifierLoc,
                                              SourceLoc &IsolatedLoc,
                                              SourceLoc &ConstLoc,
+                                             SourceLoc &DistributedLocalLoc,
                                              TypeAttributes &Attributes);
 
   bool parseConventionAttributeInternal(bool justChecking,
@@ -1335,7 +1339,8 @@ class Parser {
                                  ParamDecl::Specifier Specifier,
                                  SourceLoc SpecifierLoc,
                                  SourceLoc IsolatedLoc,
-                                 SourceLoc ConstLoc);
+                                 SourceLoc ConstLoc,
+                                 SourceLoc DistributedLocalLoc);
 
   //===--------------------------------------------------------------------===//
   // Pattern Parsing
@@ -1397,6 +1402,9 @@ class Parser {
     /// The location of the 'isolated' keyword, if present.
     SourceLoc IsolatedLoc;
 
+    /// The location of the '_local' keyword, if present.
+    SourceLoc KnownToBeLocalLoc;
+
     /// The location of the '_const' keyword, if present.
     SourceLoc CompileConstLoc;
 

include/swift/Sema/ConstraintSystem.h

@@ -1312,6 +1312,9 @@ class Solution {
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallVector<ParamDecl *, 2> isolatedParams;
 
+  /// The set of parameters that have been inferred to be '_local'.
+  llvm::SmallVector<ParamDecl *, 2> distributedKnownLocalParams;
+
   /// The set of functions that have been transformed by a result builder.
   llvm::MapVector<AnyFunctionRef, AppliedBuilderTransform>
       resultBuilderTransformed;
@@ -2674,6 +2677,10 @@ class ConstraintSystem {
   /// The set of parameters that have been inferred to be 'isolated'.
   llvm::SmallSetVector<ParamDecl *, 2> isolatedParams;
 
+  /// The set of DistributedActor parameters that have been inferred
+  /// to be known to be '_local'.
+  llvm::SmallSetVector<ParamDecl *, 2> distributedKnownLocalParams;
+
   /// Maps closure parameters to type variables.
   llvm::DenseMap<const ParamDecl *, TypeVariableType *>
     OpenedParameterTypes;
@@ -3244,6 +3251,9 @@ class ConstraintSystem {
     /// The length of \c isolatedParams.
     unsigned numIsolatedParams;
 
+    /// The length of \c distributedKnownLocalParams.
+    unsigned numDistributedKnownLocalParams;
+
     /// The length of \c ImplicitValueConversions.
     unsigned numImplicitValueConversions;
 

lib/AST/ASTDumper.cpp

@@ -753,8 +753,8 @@ namespace {
         PrintWithColorRAII(OS, DeclModifierColor) << " lazy";
       printStorageImpl(VD);
       printAccessors(VD);
-      if (VD->getAttrs().hasAttribute<KnownToBeLocalAttr>()) {
-        OS << " known-to-be-local";
+      if (VD->isDistributedKnownToBeLocal()) {
+        OS << " _local";
       }
       PrintWithColorRAII(OS, ParenthesisColor) << ')';
     }
@@ -912,8 +912,8 @@ namespace {
       OS.indent(Indent);
       PrintWithColorRAII(OS, ParenthesisColor) << '(';
       PrintWithColorRAII(OS, ParameterColor) << "parameter ";
-      if (P->getAttrs().hasAttribute<KnownToBeLocalAttr>()) {
-        OS << "known-to-be-local ";
+      if (P->isDistributedKnownToBeLocal()) {
+        OS << "_local ";
       }
       printDeclName(P);
       if (!P->getArgumentName().empty())
@@ -1349,8 +1349,8 @@ void ValueDecl::dumpRef(raw_ostream &os) const {
     os << moduleName;
   }
 
-  if (getAttrs().hasAttribute<KnownToBeLocalAttr>()) {
-    os << " known-to-be-local";
+  if (getAttrs().hasAttribute<DistributedKnownToBeLocalAttr>()) {
+    os << " _local";
   }
 
   if (getAttrs().hasAttribute<DistributedThunkAttr>()) {
@@ -3130,6 +3130,12 @@ class PrintTypeRepr : public TypeReprVisitor<PrintTypeRepr> {
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
 
+  void visitDistributedKnownToBeLocalTypeRepr(DistributedKnownToBeLocalTypeRepr *T) {
+    printCommon("_local") << '\n';
+    printRec(T->getBase());
+    PrintWithColorRAII(OS, ParenthesisColor) << ')';
+  }
+
   void visitOptionalTypeRepr(OptionalTypeRepr *T) {
     printCommon("type_optional") << '\n';
     printRec(T->getBase());
@@ -3536,6 +3542,7 @@ namespace {
       printFlag(paramFlags.isAutoClosure(), "autoclosure");
       printFlag(paramFlags.isNonEphemeral(), "nonEphemeral");
       printFlag(paramFlags.isCompileTimeConst(), "compileTimeConst");
+      printFlag(paramFlags.isDistributedKnownToBeLocal(), "distributedKnownLocal");
       switch (paramFlags.getValueOwnership()) {
       case ValueOwnership::Default: break;
       case ValueOwnership::Owned: printFlag("owned"); break;

lib/AST/ASTPrinter.cpp

@@ -3680,6 +3680,9 @@ static void printParameterFlags(ASTPrinter &printer,
 
   if (flags.isCompileTimeConst())
     printer.printKeyword("_const", options, " ");
+
+  if (flags.isDistributedKnownToBeLocal())
+    printer.printKeyword("_local", options, " ");
 }
 
 void PrintAST::visitVarDecl(VarDecl *decl) {