[Distributed] Fix too restrictive distributed witness isolation checking

include/swift/AST/Decl.h

@@ -3644,7 +3644,7 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
   bool isActor() const;
 
   /// Whether this nominal type qualifies as a distributed actor, meaning that
-  /// it is either a distributed actor.
+  /// it is either a distributed actor or a DistributedActor constrained protocol.
   bool isDistributedActor() const;
 
   /// Whether this nominal type qualifies as any actor (plain or distributed).

lib/AST/Decl.cpp

@@ -9224,7 +9224,7 @@ ActorIsolation swift::getActorIsolationOfContext(DeclContext *dc) {
       auto selfDecl = isolation.getActorInstance();
       auto actorClass = selfDecl->getType()->getReferenceStorageReferent()
           ->getClassOrBoundGenericClass();
-      // FIXME: Doesn't work properly with generics
+      // FIXME: Doesn't work properly with generics #59356
       assert(actorClass && "Bad closure actor isolation?");
       return ActorIsolation::forActorInstance(actorClass)
                 .withPreconcurrency(isolation.preconcurrency());

lib/Sema/TypeCheckProtocol.cpp

@@ -3029,14 +3029,46 @@ bool ConformanceChecker::checkActorIsolation(
   bool isDistributed = refResult.isolation.isDistributedActor() &&
       !witness->getAttrs().hasAttribute<NonisolatedAttr>();
   if (isDistributed) {
-    // If we're coming from a non-distributed requirement, then the requirement
-    // must be 'throws' to accommodate failure.
-    if (!isDistributedDecl(requirement) && !isThrowsDecl(requirement))
-      missingOptions |= MissingFlags::RequirementThrows;
-
-    if (!isDistributedDecl(witness) &&
-        (isDistributedDecl(requirement) || !missingOptions))
-      missingOptions |= MissingFlags::WitnessDistributed;
+    // Check if the protocol where the requirement originates from
+    // is a distributed actor constrained one.
+    if (cast<ProtocolDecl>(requirement->getDeclContext())->isDistributedActor()) {
+      // The requirement was declared in a DistributedActor constrained proto.
+      //
+      // This means casting up to this `P` won't "strip off" the
+      // "distributed-ness" of the type, and all call-sites will be checking
+      // distributed isolation.
+      //
+      // This means that we can actually allow these specific requirements,
+      // to be witnessed without the distributed keyword (!), but they won't be
+      // possible to be called unless:
+      // - from inside the distributed actor (self),
+      // - on a known-to-be-local distributed actor reference.
+      //
+      // This allows us to implement protocols where a local distributed actor
+      // registers "call me when something happens", and that call can be
+      // expressed as non-distributed function which we are guaranteed to be
+      // able to call, since the whenLocal will give us access to this actor as
+      // known-to-be-local, so we can invoke this method.
+
+      // If the requirement is distributed, we still need to require it on the witness though.
+      // We DO allow a non-distributed requirement to be witnessed here though!
+      if (isDistributedDecl(requirement) && !isDistributedDecl(witness))
+        missingOptions |= MissingFlags::WitnessDistributed;
+    } else {
+      // The protocol requirement comes from a normal (non-distributed actor)
+      // protocol; so the only witnesses allowed are such that we can witness
+      // them using a distributed, or nonisolated functions.
+
+      // If we're coming from a non-distributed requirement,
+      // then the requirement must be 'throws' to accommodate failure.
+      if (!isThrowsDecl(requirement))
+        missingOptions |= MissingFlags::RequirementThrows;
+
+      // If the witness is distributed, it is able to witness a requirement
+      // only if the requirement is `async throws`.
+      if (!isDistributedDecl(witness) && !missingOptions)
+        missingOptions |= MissingFlags::WitnessDistributed;
+    }
   }
 
   // If we aren't missing anything, do a Sendable check and move on.

test/Distributed/Runtime/distributed_actor_local_only_witness_and_call.swift

@@ -0,0 +1,71 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-build-swift -module-name main -Xfrontend -enable-experimental-distributed -Xfrontend -disable-availability-checking -j2 -parse-as-library -I %t %s %S/../Inputs/FakeDistributedActorSystems.swift -o %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s --color
+
+// REQUIRES: executable_test
+// REQUIRES: concurrency
+// REQUIRES: distributed
+
+// rdar://76038845
+// UNSUPPORTED: use_os_stdlib
+// UNSUPPORTED: back_deployment_runtime
+
+// FIXME(distributed): Distributed actors currently have some issues on windows, isRemote always returns false. rdar://82593574
+// UNSUPPORTED: OS=windows-msvc
+
+
+import Distributed
+import FakeDistributedActorSystems
+
+typealias DefaultDistributedActorSystem = FakeRoundtripActorSystem
+
+protocol LifecycleWatch: DistributedActor where ActorSystem == FakeRoundtripActorSystem {
+  func terminated(actor id: ID) async
+}
+
+extension LifecycleWatch {
+  func watch<T: Codable>(x: Int, _ y: T) async throws {
+    // nothing here
+    print("executed: \(#function) - x = \(x), y = \(y)")
+  }
+
+  distributed func test<T: Codable & Sendable>(x: Int, _ y: T) async throws {
+    print("executed: \(#function)")
+    try await self.watch(x: x, y)
+    print("done executed: \(#function)")
+  }
+}
+
+distributed actor Worker: LifecycleWatch {
+  func terminated(actor id: ID) async {
+    print("terminated (on \(self.id)): \(id)")
+  }
+}
+
+@main struct Main {
+  static func main() async {
+    let worker: any LifecycleWatch = Worker(actorSystem: DefaultDistributedActorSystem())
+    try! await worker.test(x: 42, "on protocol")
+
+    // CHECK: executed: test(x:_:)
+    // CHECK: executed: watch(x:_:) - x = 42, y = on protocol
+    // CHECK: done executed: test(x:_:)
+
+    // FIXME: Actor isolation getting with generics is pending implementation #59356
+    do {
+      let terminatedID = Worker.ID(parse: "<terminated-id>")
+      let __secretlyKnownToBeLocal = worker
+      await __secretlyKnownToBeLocal.terminated(actor: terminatedID)
+
+      // FIXME: Once the above fixme is solved, use this real code instead:
+      //    _ = await worker.whenLocal { __secretlyKnownToBeLocal in
+      //      let terminatedID = Worker.ID(parse: "<terminated-id>")
+      //      return await __secretlyKnownToBeLocal.terminated(actor: terminatedID)
+      //    }
+    }
+    // CHECK: terminated (on ActorAddress(address: "<unique-id>")): ActorAddress(address: "<terminated-id>")
+
+    print("OK") // CHECK: OK
+  }
+}

test/Distributed/distributed_protocol_isolation.swift

@@ -100,8 +100,9 @@ protocol StrictlyLocal {
 
   func localThrows() throws
   // expected-note@-1 2{{mark the protocol requirement 'localThrows()' 'async' to allow actor-isolated conformances}}{{22-22=async }}
-
-  // TODO: localAsync
+
+  func localAsync() async
+  // expected-note@-1 2{{mark the protocol requirement 'localAsync()' 'throws' to allow actor-isolated conformances}}
 }
 
 distributed actor Nope1_StrictlyLocal: StrictlyLocal {
@@ -111,24 +112,29 @@ distributed actor Nope1_StrictlyLocal: StrictlyLocal {
   func localThrows() throws {}
   // expected-error@-1{{distributed actor-isolated instance method 'localThrows()' cannot be used to satisfy nonisolated protocol requirement}}
   // expected-note@-2{{add 'nonisolated' to 'localThrows()' to make this instance method not isolated to the actor}}
+  func localAsync() async {}
+  // expected-error@-1{{distributed actor-isolated instance method 'localAsync()' cannot be used to satisfy nonisolated protocol requirement}}
+  // expected-note@-2{{add 'nonisolated' to 'localAsync()' to make this instance method not isolated to the actor}}
 }
 distributed actor Nope2_StrictlyLocal: StrictlyLocal {
   distributed func local() {}
   // expected-error@-1{{actor-isolated distributed instance method 'local()' cannot be used to satisfy nonisolated protocol requirement}}
   distributed func localThrows() throws {}
   // expected-error@-1{{actor-isolated distributed instance method 'localThrows()' cannot be used to satisfy nonisolated protocol requirement}}
+  distributed func localAsync() async {}
+  // expected-error@-1{{actor-isolated distributed instance method 'localAsync()' cannot be used to satisfy nonisolated protocol requirement}}
 }
 distributed actor OK_StrictlyLocal: StrictlyLocal {
   nonisolated func local() {}
   nonisolated func localThrows() throws {}
+  nonisolated func localAsync() async {}
 }
 
 protocol Server {
-  func send<Message: Codable>(message: Message) async throws -> String
+  func send<Message: Codable & Sendable>(message: Message) async throws -> String
 }
 actor MyServer : Server {
-  // expected-note@+1{{consider making generic parameter 'Message' conform to the 'Sendable' protocol}} {{29-29=, Sendable}}
-  func send<Message: Codable>(message: Message) throws -> String { "" }  // expected-warning{{non-sendable type 'Message' in parameter of actor-isolated instance method 'send(message:)' satisfying protocol requirement cannot cross actor boundary}}
+  func send<Message: Codable & Sendable>(message: Message) throws -> String { "" } // OK
 }
 
 protocol AsyncThrowsAll {
@@ -140,9 +146,15 @@ actor LocalOK_AsyncThrowsAll: AsyncThrowsAll {
   func maybe(param: String, int: Int) async throws -> Int { 1111 }
 }
 
-actor LocalOK_Implicitly_AsyncThrowsAll: AsyncThrowsAll {
+actor LocalOK_ImplicitlyThrows_AsyncThrowsAll: AsyncThrowsAll {
+  func maybe(param: String, int: Int) async -> Int { 1111 }
+}
+actor LocalOK_ImplicitlyAsync_AsyncThrowsAll: AsyncThrowsAll {
   func maybe(param: String, int: Int) throws -> Int { 1111 }
 }
+actor LocalOK_ImplicitlyThrowsAsync_AsyncThrowsAll: AsyncThrowsAll {
+  func maybe(param: String, int: Int) -> Int { 1111 }
+}
 
 distributed actor Nope1_AsyncThrowsAll: AsyncThrowsAll {
   func maybe(param: String, int: Int) async throws -> Int { 111 }
@@ -170,6 +182,75 @@ func testAsyncThrowsAll(p: AsyncThrowsAll,
   _ = try await pp.maybe(param: "", int: 0)
 }
 
+// ==== -----------------------------------------------------------------------
+// MARK: Distributed actor protocols can have non-dist requirements
+
+protocol TerminationWatchingA {
+  func terminated(a: String) async
+  // expected-note@-1{{mark the protocol requirement 'terminated(a:)' 'throws' to allow actor-isolated conformances}}
+}
+
+protocol TerminationWatchingDA: DistributedActor {
+  func terminated(da: String) async // expected-note 3 {{distributed actor-isolated instance method 'terminated(da:)' declared here}}
+}
+
+actor A_TerminationWatchingA: TerminationWatchingA {
+  func terminated(a: String) { } // ok, since: actor -> implicitly async
+}
+func test_watching_A(a: A_TerminationWatchingA) async throws {
+  await a.terminated(a: "normal")
+}
+
+distributed actor DA_TerminationWatchingA: TerminationWatchingA {
+  func terminated(a: String) { }
+  // expected-error@-1{{distributed actor-isolated instance method 'terminated(a:)' cannot be used to satisfy nonisolated protocol requirement}}
+  // expected-note@-2{{add 'nonisolated' to 'terminated(a:)' to make this instance method not isolated to the actor}}
+}
+
+distributed actor DA_TerminationWatchingDA: TerminationWatchingDA {
+  distributed func test() {}
+  func terminated(da: String) { }
+  // expected-note@-1{{distributed actor-isolated instance method 'terminated(da:)' declared here}}
+}
+
+func test_watchingDA(da: DA_TerminationWatchingDA) async throws {
+  try await da.test() // ok
+  da.terminated(da: "the terminated func is not distributed") // expected-error{{only 'distributed' instance methods can be called on a potentially remote distributed actor}}
+}
+
+func test_watchingDA<WDA: TerminationWatchingDA>(da: WDA) async throws {
+  try await da.terminated(da: "the terminated func is not distributed")
+  // expected-error@-1{{only 'distributed' instance methods can be called on a potentially remote distributed actor}}
+  // expected-warning@-2{{no calls to throwing functions occur within 'try' expression}}
+
+  let __secretlyKnownToBeLocal = da
+  await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK // FIXME(#59356): (the __secretlyKnown is a hack, but the whenLocal crashes now on pending isolation getting with generic actors for closures)
+  // FIXME: pending fix of closure isolation checking with actors #59356
+  // await da.whenLocal { __secretlyKnownToBeLocal in
+  //   await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
+  // }
+}
+
+func test_watchingDA_erased(da: DA_TerminationWatchingDA) async throws {
+  let wda: any TerminationWatchingDA = da
+  try await wda.terminated(da: "the terminated func is not distributed")
+  // expected-error@-1{{only 'distributed' instance methods can be called on a potentially remote distributed actor}}
+  // expected-warning@-2{{no calls to throwing functions occur within 'try' expression}}
+
+  let __secretlyKnownToBeLocal = wda
+  await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK // FIXME(#59356): (the __secretlyKnown is a hack, but the whenLocal crashes now on pending isolation getting with generic actors for closures)
+  // FIXME: pending fix of closure isolation checking with actors #59356
+  // await wda.whenLocal { __secretlyKnownToBeLocal in
+  //   await __secretlyKnownToBeLocal.terminated(da: "local calls are okey!") // OK
+  // }
+}
+
+func test_watchingDA_any(da: any TerminationWatchingDA) async throws {
+  try await da.terminated(da: "the terminated func is not distributed")
+  // expected-error@-1{{only 'distributed' instance methods can be called on a potentially remote distributed actor}}
+  // expected-warning@-2{{no calls to throwing functions occur within 'try' expression}}
+}
+
 // ==== ------------------------------------------------------------------------
 // MARK: Error cases