[Distributed] Handle distributed func witnessess from distributed actor protocols (#38269)

* [Distributed] Handle distributed func witnessess from distributed actor protocols

* [Distributed] Implement dist protocol and nonisolated handling

* revert additional warning check

* [Distributed] Ban mixing nonisolated and distributed on func

* [Distributed] handle nonisolated in distributed contexts

Author: ktoso

include/swift/AST/Decl.h

@@ -3280,6 +3280,9 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   /// it is either a distributed actor.
   bool isDistributedActor() const;
 
+  /// Whether this nominal type qualifies as any actor (plain or distributed).
+  bool isAnyActor() const;
+
   /// Return the range of semantics attributes attached to this NominalTypeDecl.
   auto getSemanticsAttrs() const
       -> decltype(getAttrs().getSemanticsAttrs()) {

include/swift/AST/DiagnosticsSema.def

@@ -4441,7 +4441,7 @@ NOTE(distributed_actor_isolated_method_note,none,
       "only 'distributed' functions can be called from outside the distributed actor",
       ())
 ERROR(distributed_actor_isolated_method,none,
-      "only 'distributed' functions can be called from outside the distributed actor", // TODO: more like 'non-distributed' ... defined here
+      "only 'distributed' functions can be called from outside the distributed actor", // TODO: improve error message to be more like 'non-distributed' ... defined here
       ())
 ERROR(distributed_actor_func_param_not_codable,none,
       "distributed function parameter '%0' of type %1 does not conform to 'Codable'",
@@ -4452,6 +4452,12 @@ ERROR(distributed_actor_func_result_not_codable,none,
 ERROR(distributed_actor_remote_func_implemented_manually,none,
       "Distributed function's %0 remote counterpart %1 cannot not be implemented manually.",
       (Identifier, Identifier))
+ERROR(nonisolated_distributed_actor_storage,none,
+      "'nonisolated' can not be applied to distributed actor stored properties",
+      ())
+ERROR(distributed_actor_func_nonisolated, none,
+      "function %0 cannot be both 'nonisolated' and 'distributed'",
+      (DeclName))
 ERROR(distributed_actor_remote_func_is_not_static,none,
       "remote function %0 must be static.",
       (DeclName))
@@ -4600,7 +4606,7 @@ ERROR(distributed_actor_independent_property_must_be_let,none,
       "_distributedActorIndependent can be applied to properties, however they must be 'let'",
       ())
 NOTE(distributed_actor_isolated_property,none,
-      "distributed actor state is only available within the actor instance",
+      "distributed actor state is only available within the actor instance", // TODO: reword in terms of isolation
       ())
 
 ERROR(concurrency_lib_missing,none,

lib/AST/ASTDumper.cpp

@@ -855,6 +855,10 @@ namespace {
         D->getCaptureInfo().print(OS);
       }
 
+      if (D->isDistributed()) {
+        OS << " distributed";
+      }
+
       if (auto fac = D->getForeignAsyncConvention()) {
         OS << " foreign_async=";
         if (auto type = fac->completionHandlerType())

lib/AST/Decl.cpp

@@ -3895,6 +3895,10 @@ bool NominalTypeDecl::isDistributedActor() const {
                            false);
 }
 
+bool NominalTypeDecl::isAnyActor() const {
+  return isActor() || isDistributedActor();
+}
+
 GenericTypeDecl::GenericTypeDecl(DeclKind K, DeclContext *DC,
                                  Identifier name, SourceLoc nameLoc,
                                  ArrayRef<InheritedEntry> inherited,

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -527,7 +527,6 @@ static void addImplicitResignAddress(ClassDecl *decl) {
   BraceStmt *newBody = BraceStmt::create(C, SourceLoc(), statements, SourceLoc(),
                                          /*implicit=*/true);
 
-  newBody->dump();
   deinitDecl->setBody(newBody, AbstractFunctionDecl::BodyKind::TypeChecked); // FIXME: no idea if Parsed is right, we are NOT type checked I guess?
 }
 

lib/Sema/TypeCheckAttr.cpp

@@ -5426,7 +5426,8 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
       return;
     }
   } else if (dyn_cast<StructDecl>(D) || dyn_cast<EnumDecl>(D)) {
-    diagnoseAndRemoveAttr(attr, diag::distributed_actor_func_not_in_distributed_actor);
+    diagnoseAndRemoveAttr(
+        attr, diag::distributed_actor_func_not_in_distributed_actor);
     return;
   }
 
@@ -5437,19 +5438,31 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
       return;
     }
 
+    // distributed func cannot be simultaneously nonisolated
+    if (auto nonisolated =
+            funcDecl->getAttrs().getAttribute<NonisolatedAttr>()) {
+      diagnoseAndRemoveAttr(nonisolated,
+                            diag::distributed_actor_func_nonisolated,
+                            funcDecl->getName());
+      return;
+    }
+
     // distributed func must be declared inside an distributed actor
     if (dc->getSelfClassDecl() &&
         !dc->getSelfClassDecl()->isDistributedActor()) {
-      diagnoseAndRemoveAttr(attr, diag::distributed_actor_func_not_in_distributed_actor);
+      diagnoseAndRemoveAttr(
+          attr, diag::distributed_actor_func_not_in_distributed_actor);
       return;
     } else if (auto protoDecl = dc->getSelfProtocolDecl()){
       if (!protoDecl->inheritsFromDistributedActor()) {
         // TODO: could suggest adding `: DistributedActor` to the protocol as well
-        diagnoseAndRemoveAttr(attr, diag::distributed_actor_func_not_in_distributed_actor);
+        diagnoseAndRemoveAttr(
+            attr, diag::distributed_actor_func_not_in_distributed_actor);
         return;
       }
     } else if (dc->getSelfStructDecl() || dc->getSelfEnumDecl()) {
-      diagnoseAndRemoveAttr(attr, diag::distributed_actor_func_not_in_distributed_actor);
+      diagnoseAndRemoveAttr(
+          attr, diag::distributed_actor_func_not_in_distributed_actor);
       return;
     }
   }
@@ -5459,11 +5472,27 @@ void AttributeChecker::visitNonisolatedAttr(NonisolatedAttr *attr) {
   // 'nonisolated' can be applied to global and static/class variables
   // that do not have storage.
   auto dc = D->getDeclContext();
+
   if (auto var = dyn_cast<VarDecl>(D)) {
-    // 'nonisolated' can not be applied to mutable stored properties.
-    if (var->hasStorage() && var->supportsMutation()) {
-      diagnoseAndRemoveAttr(attr, diag::nonisolated_mutable_storage);
-      return;
+    // stored properties have limitations as to when they can be nonisolated.
+    if (var->hasStorage()) {
+      auto nominal = dyn_cast<NominalTypeDecl>(dc);
+
+      // 'nonisolated' can not be applied to stored properties inside
+      // distributed actors. Attempts of nonisolated access would be
+      // cross-actor, and that means they might be accessing on a remote actor,
+      // in which case the stored property storage does not exist.
+      if (nominal && nominal->isDistributedActor()) {
+        diagnoseAndRemoveAttr(attr,
+                              diag::nonisolated_distributed_actor_storage);
+        return;
+      }
+
+      // 'nonisolated' can not be applied to mutable stored properties.
+      if (var->supportsMutation()) {
+        diagnoseAndRemoveAttr(attr, diag::nonisolated_mutable_storage);
+        return;
+      }
     }
 
     // nonisolated can not be applied to local properties.

lib/Sema/TypeCheckConcurrency.cpp

@@ -393,10 +393,10 @@ ActorIsolationRestriction ActorIsolationRestriction::forDeclaration(
                 diag::distributed_actor_func_defined_outside_of_distributed_actor,
                 func->getName());
           }
-        }
+        } // TODO: need to handle protocol case here too?
 
-          return forDistributedActorSelf(isolation.getActor(),
-                                         /*isCrossActor*/ isAccessibleAcrossActors); // TODO: not sure?
+        return forDistributedActorSelf(isolation.getActor(),
+                                       /*isCrossActor*/ isAccessibleAcrossActors);
       }
     }
 
@@ -438,7 +438,6 @@ ActorIsolationRestriction ActorIsolationRestriction::forDeclaration(
     }
 
     case ActorIsolation::Independent:
-      // Actor-independent have no restrictions on their access.
       return forUnrestricted();
 
     case ActorIsolation::Unspecified:
@@ -2264,8 +2263,27 @@ namespace {
       switch (auto isolation =
                   ActorIsolationRestriction::forDeclaration(
                     memberRef, getDeclContext())) {
-      case ActorIsolationRestriction::Unrestricted:
+      case ActorIsolationRestriction::Unrestricted: {
+        // If a cross-actor reference is to an isolated actor, it's not
+        // crossing actors.
+        if (getIsolatedActor(base))
+          return false;
+
+        // Always fine to invoke constructors from outside of actors.
+        if (dyn_cast<ConstructorDecl>(member))
+          return false;
+
+        // While the member may be unrestricted, perhaps
+        if (auto classDecl = dyn_cast<ClassDecl>(member->getDeclContext())) {
+          if (classDecl->isDistributedActor()) {
+            ctx.Diags.diagnose(memberLoc, diag::distributed_actor_isolated_method);
+            noteIsolatedActorMember(member, context);
+            return true;
+          }
+        }
+
         return false;
+      }
 
       case ActorIsolationRestriction::CrossActorSelf: {
         // If a cross-actor reference is to an isolated actor, it's not
@@ -2295,7 +2313,6 @@ namespace {
               member->isInstanceMember() &&
               isActorInitOrDeInitContext(getDeclContext()))) {
           // invocation on not-'self', is only okey if this is a distributed func
-
           if (auto func = dyn_cast<FuncDecl>(member)) {
             if (!func->isDistributed()) {
               ctx.Diags.diagnose(memberLoc, diag::distributed_actor_isolated_method);
@@ -2411,6 +2428,14 @@ namespace {
       case ActorIsolationRestriction::Unsafe:
         // This case is hit when passing actor state inout to functions in some
         // cases. The error is emitted by diagnoseInOutArg.
+
+        if (auto classDecl = dyn_cast<ClassDecl>(member->getDeclContext())) {
+          if (classDecl->isDistributedActor()) {
+            member->diagnose(diag::distributed_actor_isolated_method);
+            return true;
+          }
+        }
+
         return false;
       }
       llvm_unreachable("unhandled actor isolation kind!");
@@ -3132,11 +3157,17 @@ ActorIsolation ActorIsolationRequest::evaluate(
   // overridden by other inference rules.
   ActorIsolation defaultIsolation = ActorIsolation::forUnspecified();
 
-  // A @Sendable function is assumed to be actor-independent.
   if (auto func = dyn_cast<AbstractFunctionDecl>(value)) {
+    // A @Sendable function is assumed to be actor-independent.
     if (func->isSendable()) {
       defaultIsolation = ActorIsolation::forIndependent();
     }
+
+    if (auto nominal = value->getDeclContext()->getSelfNominalTypeDecl()) {
+      if (nominal->isDistributedActor()) {
+        defaultIsolation = ActorIsolation::forDistributedActorInstance(nominal);
+      }
+    }
   }
 
   // An actor's convenience init is assumed to be actor-independent.
@@ -3172,10 +3203,14 @@ ActorIsolation ActorIsolationRequest::evaluate(
     }
 
     case ActorIsolation::DistributedActorInstance: {
-      /// 'distributed actor independent' implies 'actor independent'
-      if (value->isDistributedActorIndependent())
+      /// 'distributed actor independent' implies 'nonisolated'
+      if (value->isDistributedActorIndependent()) {
+        // TODO: rename 'distributed actor independent' to 'distributed(nonisolated)'
         value->getAttrs().add(
             new (ctx) DistributedActorIndependentAttr(/*IsImplicit=*/true));
+        value->getAttrs().add(
+            new (ctx) NonisolatedAttr(/*IsImplicit=*/true));
+      }
       break;
     }
     case ActorIsolation::ActorInstance:
@@ -3186,6 +3221,7 @@ ActorIsolation ActorIsolationRequest::evaluate(
       // Nothing to do.
       break;
     }
+
     return inferred;
   };
 
@@ -3291,7 +3327,7 @@ bool HasIsolatedSelfRequest::evaluate(
   // Only ever applies to members of actors.
   auto dc = value->getDeclContext();
   auto selfTypeDecl = dc->getSelfNominalTypeDecl();
-  if (!selfTypeDecl || !selfTypeDecl->isActor())
+  if (!selfTypeDecl || !selfTypeDecl->isAnyActor())
     return false;
 
   // For accessors, consider the storage declaration.

lib/Sema/TypeCheckProtocol.cpp

@@ -2827,31 +2827,28 @@ bool ConformanceChecker::checkActorIsolation(
                   /*fromExpression=*/false)) {
   case ActorIsolationRestriction::DistributedActorSelf: {
     if (witness->isSynthesized()) {
-      // Some of our synthesized properties get special treatment,
-      // they are always available, regardless if the actor is remote even.
-      auto &C = requirement->getASTContext();
-
-      // actorAddress is special, it is *always* available.
-      // even if the actor is 'remote' it is always available and immutable.
-      if (witness->getName() == C.Id_actorAddress &&
-          witness->getInterfaceType()->isEqual(
-              C.getActorAddressDecl()->getDeclaredInterfaceType()))
+      // we only have two 'distributed-actor-nonisolated' properties,
+      // the address and transport; if we see any such marked property,
+      // we're free to automatically assume those are fine and accessible always.
+      if (witness->isDistributedActorIndependent())
         return false;
+    }
 
-      // TODO: we don't *really* need to expose the transport like that... reconsider?
-      if (witness->getName() == C.Id_actorTransport &&
-          witness->getInterfaceType()->isEqual(
-              C.getActorTransportDecl()->getDeclaredInterfaceType()))
+    if (auto funcDecl = dyn_cast<AbstractFunctionDecl>(witness)) {
+      // A 'distributed func' may witness a distributed isolated function requirement.
+      if (funcDecl->isDistributed())
         return false;
     }
 
     // continue checking ActorSelf rules
     LLVM_FALLTHROUGH;
   }
   case ActorIsolationRestriction::ActorSelf: {
+    auto requirementIsolation = getActorIsolation(requirement);
+
     // An actor-isolated witness can only conform to an actor-isolated
     // requirement.
-    if (getActorIsolation(requirement) == ActorIsolation::ActorInstance)
+    if (requirementIsolation == ActorIsolation::ActorInstance)
       return false;
 
     witness->diagnose(diag::actor_isolated_witness,
@@ -2903,8 +2900,12 @@ bool ConformanceChecker::checkActorIsolation(
   bool requirementIsUnsafe = false;
   switch (auto requirementIsolation = getActorIsolation(requirement)) {
   case ActorIsolation::ActorInstance:
-  case ActorIsolation::DistributedActorInstance:
     llvm_unreachable("There are not actor protocols");
+  case ActorIsolation::DistributedActorInstance:
+    // A requirement inside a distributed actor, where it has a protocol that was
+    // bound requiring a `DistributedActor` conformance (`protocol D: DistributedActor`),
+    // results in the requirement being isolated to given distributed actor.
+    break;
 
   case ActorIsolation::GlobalActorUnsafe:
     requirementIsUnsafe = true;

stdlib/public/Distributed/DistributedActor.swift

@@ -97,6 +97,7 @@ extension DistributedActor {
         fatalError("\(#function) is not implemented yet for distributed actors'")
     }
 
+    // FIXME: distributed(nonisolated)
     nonisolated public func encode(to encoder: Encoder) throws {
         var container = encoder.singleValueContainer()
         try container.encode(self.actorAddress)

test/Distributed/distributed_actor_basic.swift

@@ -0,0 +1,30 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-distributed
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+import _Distributed
+
+@available(SwiftStdlib 5.5, *)
+distributed struct StructNope {} // expected-error{{distributed' modifier cannot be applied to this declaration}}
+@available(SwiftStdlib 5.5, *)
+distributed class ClassNope {} // expected-error{{'distributed' can only be applied to 'actor' definitions, and distributed actor-isolated async functions}}
+@available(SwiftStdlib 5.5, *)
+distributed enum EnumNope {} // expected-error{{distributed' modifier cannot be applied to this declaration}}
+
+@available(SwiftStdlib 5.5, *)
+distributed actor DA {
+
+  class func classFunc() {
+    // expected-error@-1{{class methods are only allowed within classes; use 'static' to declare a static method}}
+  }
+
+  nonisolated distributed func nonisolatedDistributed() async {
+    // expected-error@-1{{function 'nonisolatedDistributed()' cannot be both 'nonisolated' and 'distributed'}}{{3-15=}}
+    fatalError()
+  }
+
+  distributed nonisolated func distributedNonisolated() async {
+    // expected-error@-1{{function 'distributedNonisolated()' cannot be both 'nonisolated' and 'distributed'}}{{15-27=}}
+    fatalError()
+  }
+}

test/Distributed/distributed_actor_initialization.swift

@@ -156,5 +156,3 @@ distributed actor BadRedeclare22 { //expected-error {{type 'BadRedeclare22' does
   // expected-error@-3 {{invalid redeclaration of synthesized 'init(resolve:using:)'}}
   // expected-error@-4 {{invalid redeclaration of synthesized initializer 'init(resolve:using:)'}}
 }
-
-// TODO: handle subclassing as well

test/Distributed/distributed_actor_isolation.swift

@@ -16,13 +16,6 @@ actor LocalActor_1 {
 
 struct NotCodableValue { }
 
-@available(SwiftStdlib 5.5, *)
-distributed struct StructNope {} // expected-error{{distributed' modifier cannot be applied to this declaration}}
-@available(SwiftStdlib 5.5, *)
-distributed class ClassNope {} // expected-error{{'distributed' can only be applied to 'actor' definitions, and distributed actor-isolated async functions}}
-@available(SwiftStdlib 5.5, *)
-distributed enum EnumNope {} // expected-error{{distributed' modifier cannot be applied to this declaration}}
-
 @available(SwiftStdlib 5.5, *)
 distributed actor DistributedActor_1 {