[Distributed] Implement missing case in permitting witnesses

Author: ktoso

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/Sema/TypeCheckProtocol.cpp

@@ -3029,14 +3029,47 @@ 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.
+    auto proto = dyn_cast<ProtocolDecl>(requirement->getDeclContext());
+    if (proto && proto->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 (!isDistributedDecl(requirement) && !isThrowsDecl(requirement))
+        missingOptions |= MissingFlags::RequirementThrows;
+
+      // If the requirement is distributed, we require a distributed witness
+      if (!isDistributedDecl(witness) &&
+          (isDistributedDecl(requirement) || !missingOptions))
+        missingOptions |= MissingFlags::WitnessDistributed;
+    }
   }
 
   // If we aren't missing anything, do a Sendable check and move on.

test/Distributed/distributed_protocol_isolation.swift

@@ -170,6 +170,56 @@ 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@-1{{distributed actor-isolated instance method 'terminated(da:)' declared here}}
+  // expected-note@-2{{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}}
+}
+
+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