[1/3][Distributed] Make distributed thunks the witnesses, fix calls on generic DAs

include/swift/AST/DiagnosticsSema.def

@@ -4791,6 +4791,9 @@ ERROR(distributed_actor_func_static,none,
 ERROR(distributed_actor_func_not_in_distributed_actor,none,
       "'distributed' method can only be declared within 'distributed actor'",
       ())
+ERROR(distributed_method_requirement_must_be_async_throws,none, // FIXME(distributed): this is an implementation limitation we should lift
+      "'distributed' protocol requirement %0 must currently be declared explicitly 'async throws'",
+      (DeclName))
 ERROR(distributed_actor_user_defined_special_property,none,
       "property %0 cannot be defined explicitly, as it conflicts with "
       "distributed actor synthesized stored property",

include/swift/AST/DistributedDecl.h

@@ -47,6 +47,11 @@ Type getDistributedActorSystemType(NominalTypeDecl *actor);
 /// Determine the `ID` type for the given actor.
 Type getDistributedActorIDType(NominalTypeDecl *actor);
 
+/// Similar to `getDistributedSerializationRequirementType`, however, from the
+/// perspective of a concrete function. This way we're able to get the
+/// serialization requirement for specific members, also in protocols.
+Type getConcreteReplacementForMemberSerializationRequirement(ValueDecl *member);
+
 /// Get specific 'SerializationRequirement' as defined in 'nominal'
 /// type, which must conform to the passed 'protocol' which is expected
 /// to require the 'SerializationRequirement'.

include/swift/SIL/SILDeclRef.h

@@ -414,8 +414,7 @@ struct SILDeclRef {
                       defaultArgIndex,
                       pointer.get<AutoDiffDerivativeFunctionIdentifier *>());
   }
-  /// Returns the distributed entry point corresponding to the same
-  /// decl.
+  /// Returns the distributed entry point corresponding to the same decl.
   SILDeclRef asDistributed(bool distributed = true) const {
     return SILDeclRef(loc.getOpaqueValue(), kind,
                       /*foreign=*/false,

lib/AST/DistributedDecl.cpp

@@ -96,6 +96,37 @@ Type swift::getConcreteReplacementForProtocolActorSystemType(ValueDecl *member)
   llvm_unreachable("Unable to fetch ActorSystem type!");
 }
 
+Type swift::getConcreteReplacementForMemberSerializationRequirement(
+    ValueDecl *member) {
+  auto &C = member->getASTContext();
+  auto *DC = member->getDeclContext();
+  auto DA = C.getDistributedActorDecl();
+
+  // === When declared inside an actor, we can get the type directly
+  if (auto classDecl = DC->getSelfClassDecl()) {
+    return getDistributedSerializationRequirementType(classDecl, C.getDistributedActorDecl());
+  }
+
+  /// === Maybe the value is declared in a protocol?
+  if (auto protocol = DC->getSelfProtocolDecl()) {
+    GenericSignature signature;
+    if (auto *genericContext = member->getAsGenericContext()) {
+      signature = genericContext->getGenericSignature();
+    } else {
+      signature = DC->getGenericSignatureOfContext();
+    }
+
+    auto SerReqAssocType = DA->getAssociatedType(C.Id_SerializationRequirement)
+                               ->getDeclaredInterfaceType();
+
+    // Note that this may be null, e.g. if we're a distributed func inside
+    // a protocol that did not declare a specific actor system requirement.
+    return signature->getConcreteType(SerReqAssocType);
+  }
+
+  llvm_unreachable("Unable to fetch ActorSystem type!");
+}
+
 Type swift::getDistributedActorSystemType(NominalTypeDecl *actor) {
   assert(!dyn_cast<ProtocolDecl>(actor) &&
          "Use getConcreteReplacementForProtocolActorSystemType instead to get"

lib/SILGen/SILGenPoly.cpp

@@ -4439,15 +4439,22 @@ void SILGenFunction::emitProtocolWitness(
   SmallVector<ManagedValue, 8> origParams;
   collectThunkParams(loc, origParams);
 
-  // If the witness is isolated to a distributed actor, but the requirement is
-  // not, go through the distributed thunk.
   if (witness.hasDecl() &&
-      getActorIsolation(witness.getDecl()).isDistributedActor() &&
-      requirement.hasDecl() &&
-      !getActorIsolation(requirement.getDecl()).isDistributedActor()) {
-    witness = SILDeclRef(
-        cast<AbstractFunctionDecl>(witness.getDecl())->getDistributedThunk())
-          .asDistributed();
+      getActorIsolation(witness.getDecl()).isDistributedActor()) {
+    // We witness protocol requirements using the distributed thunk, when:
+    // - the witness is isolated to a distributed actor, but the requirement is not
+    // - the requirement is a distributed func, and therefore can only be witnessed
+    //   by a distributed func; we handle this by witnessing the requirement with the thunk
+    // FIXME(distributed): this limits us to only allow distributed explicitly throwing async requirements... we need to fix this somehow.
+    if (requirement.hasDecl()) {
+      if ((!getActorIsolation(requirement.getDecl()).isDistributedActor()) ||
+          (isa<FuncDecl>(requirement.getDecl()) &&
+              witness.getFuncDecl()->isDistributed())) {
+        auto thunk = cast<AbstractFunctionDecl>(witness.getDecl())
+                         ->getDistributedThunk();
+        witness = SILDeclRef(thunk).asDistributed();
+      }
+    }
   } else if (enterIsolation) {
     // If we are supposed to enter the actor, do so now by hopping to the
     // actor.

lib/Sema/CodeSynthesisDistributedActor.cpp

@@ -648,8 +648,13 @@ static FuncDecl *createDistributedThunkFunction(FuncDecl *func) {
   auto &C = func->getASTContext();
   auto DC = func->getDeclContext();
 
-  auto systemTy = getConcreteReplacementForProtocolActorSystemType(func);
-  assert(systemTy &&
+  // NOTE: So we don't need a thunk in the protocol, we should call the underlying
+  // thing instead, which MUST have a thunk, since it must be a distributed func as well...
+  if (dyn_cast<ProtocolDecl>(DC)) {
+    return nullptr;
+  }
+
+  assert(getConcreteReplacementForProtocolActorSystemType(func) &&
          "Thunk synthesis must have concrete actor system type available");
 
   DeclName thunkName = func->getName();

lib/Sema/TypeCheckAttr.cpp

@@ -5900,6 +5900,7 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
 
     // distributed func must be declared inside an distributed actor
     auto selfTy = dc->getSelfTypeInContext();
+
     if (!selfTy->isDistributedActor()) {
       auto diagnostic = diagnoseAndRemoveAttr(
         attr, diag::distributed_actor_func_not_in_distributed_actor);
@@ -5910,6 +5911,24 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
       }
       return;
     }
+
+    // Diagnose for the limitation that we currently have to require distributed
+    // actor constrained protocols to declare the distributed requirements as
+    // 'async throws'
+    // FIXME: rdar://95949498 allow requirements to not declare explicit async/throws in protocols; those effects are implicit in any case
+    if (isa<ProtocolDecl>(dc)) {
+      if (!funcDecl->hasAsync() || !funcDecl->hasThrows()) {
+        auto diag = funcDecl->diagnose(diag::distributed_method_requirement_must_be_async_throws,
+                           funcDecl->getName());
+        if (!funcDecl->hasAsync()) {
+          diag.fixItInsertAfter(funcDecl->getThrowsLoc(), " async");
+        }
+        if (!funcDecl->hasThrows()) {
+          diag.fixItInsertAfter(funcDecl->getThrowsLoc(), " throws");
+        }
+        return;
+      }
+    }
   }
 }
 

lib/Sema/TypeCheckDistributed.cpp

@@ -503,9 +503,25 @@ bool CheckDistributedFunctionRequest::evaluate(
     serializationRequirements = getDistributedSerializationRequirementProtocols(
         getDistributedActorSystemType(actor)->getAnyNominal(),
         C.getProtocol(KnownProtocolKind::DistributedActorSystem));
+  } else if (isa<ProtocolDecl>(DC)) {
+    if (auto seqReqTy =
+        getConcreteReplacementForMemberSerializationRequirement(func)) {
+      auto seqReqTyDes = seqReqTy->castTo<ExistentialType>()->getConstraintType()->getDesugaredType();
+      for (auto req : flattenDistributedSerializationTypeToRequiredProtocols(seqReqTyDes)) {
+        serializationRequirements.insert(req);
+      }
+    }
+
+    // The distributed actor constrained protocol has no serialization requirements
+    // or actor system defined, so these will only be enforced, by implementations
+    // of DAs conforming to it, skip checks here.
+    if (serializationRequirements.empty()) {
+      return false;
+    }
   } else {
-    llvm_unreachable("Cannot handle types other than extensions and actor "
-                     "declarations in distributed function checking.");
+    llvm_unreachable("Distributed function detected in type other than extension, "
+                     "distributed actor, or protocol! This should not be possible "
+                     ", please file a bug.");
   }
 
   // If the requirement is exactly `Codable` we diagnose it ia bit nicer.
@@ -653,12 +669,23 @@ void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal
   // If applicable, this will create the default 'init(transport:)' initializer
   (void)nominal->getDefaultInitializer();
 
+
   for (auto member : nominal->getMembers()) {
     // --- Ensure all thunks
     if (auto func = dyn_cast<AbstractFunctionDecl>(member)) {
       if (!func->isDistributed())
         continue;
 
+      if (!isa<ProtocolDecl>(nominal)) {
+        auto systemTy = getConcreteReplacementForProtocolActorSystemType(func);
+        if (!systemTy || systemTy->hasError()) {
+          nominal->diagnose(
+              diag::distributed_actor_conformance_missing_system_type,
+              nominal->getName());
+          return;
+        }
+      }
+
       if (auto thunk = func->getDistributedThunk()) {
         SF->DelayedFunctions.push_back(thunk);
       }

test/Distributed/Runtime/distributed_actor_func_calls_remoteCall_through_generic.swift

@@ -0,0 +1,165 @@
+// 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 -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 --dump-input=always
+
+// 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 DistributedWorker: DistributedActor where ActorSystem == DefaultDistributedActorSystem {
+  associatedtype WorkItem: Sendable & Codable
+  associatedtype WorkResult: Sendable & Codable
+
+  // distributed requirement currently is forced to be `async throws`...
+  // FIXME(distributed): requirements don't have to be async throws,
+  //                     distributed makes them implicitly async throws anyway...
+  distributed func submit(work: WorkItem) async throws -> WorkResult
+
+  // non distributed requirements can be witnessed with _normal_ functions
+  func sync(work: WorkItem) -> WorkResult
+  func async(work: WorkItem) async -> WorkResult
+  func syncThrows(work: WorkItem) throws -> WorkResult
+  func asyncThrows(work: WorkItem) async throws -> WorkResult
+}
+
+distributed actor TheWorker: DistributedWorker {
+  typealias ActorSystem = DefaultDistributedActorSystem
+  typealias WorkItem = String
+  typealias WorkResult = String
+
+  distributed func submit(work: WorkItem) async throws -> WorkResult {
+    "\(#function): \(work)"
+  }
+
+  func sync(work: WorkItem) -> WorkResult {
+    return "\(#function): \(work)"
+  }
+  func async(work: WorkItem) async -> WorkResult {
+    return "\(#function): \(work)"
+  }
+  func syncThrows(work: WorkItem) throws -> WorkResult {
+    return "\(#function): \(work)"
+  }
+  func asyncThrows(work: WorkItem) async throws -> WorkResult {
+    return "\(#function): \(work)"
+  }
+}
+
+func test_generic(system: DefaultDistributedActorSystem) async throws {
+  let localW = TheWorker(actorSystem: system)
+  let remoteW = try! TheWorker.resolve(id: localW.id, using: system)
+  precondition(__isRemoteActor(remoteW))
+
+  // direct calls work ok:
+  let replyDirect = try await remoteW.submit(work: "Direct")
+  print("reply direct: \(replyDirect)")
+  // CHECK: >> remoteCall: on:main.TheWorker, target:main.TheWorker.submit(work:), invocation:FakeInvocationEncoder(genericSubs: [], arguments: ["Direct"], returnType: Optional(Swift.String), errorType: Optional(Swift.Error)), throwing:Swift.Error, returning:Swift.String
+  // CHECK: reply direct: submit(work:): Direct
+
+  func callWorker<W: DistributedWorker>(w: W) async throws -> String where W.WorkItem == String, W.WorkResult == String {
+    try await w.submit(work: "Hello")
+  }
+  let reply = try await callWorker(w: remoteW)
+  print("reply (remote): \(reply)")
+  // CHECK: >> remoteCall: on:main.TheWorker, target:main.TheWorker.submit(work:), invocation:FakeInvocationEncoder(genericSubs: [], arguments: ["Hello"], returnType: Optional(Swift.String), errorType: Optional(Swift.Error)), throwing:Swift.Error, returning:Swift.String
+  // CHECK: << remoteCall return: submit(work:): Hello
+  // CHECK: reply (remote): submit(work:): Hello
+
+  let replyLocal = try await callWorker(w: localW)
+  print("reply (local): \(replyLocal)")
+  // CHECK-NOT: >> remoteCall
+  // CHECK: reply (local): submit(work:): Hello
+}
+
+func test_whenLocal(system: DefaultDistributedActorSystem) async throws {
+  let localW = TheWorker(actorSystem: system)
+  let remoteW = try! TheWorker.resolve(id: localW.id, using: system)
+  precondition(__isRemoteActor(remoteW))
+
+  do {
+    let replySync = await remoteW.whenLocal { __secretlyKnownToBeLocal in
+      __secretlyKnownToBeLocal.sync(work: "test")
+    }
+    print("replySync (remote): \(replySync)")
+    // CHECK: replySync (remote): nil
+
+    let replySyncThrows = try await remoteW.whenLocal { __secretlyKnownToBeLocal in
+      try __secretlyKnownToBeLocal.syncThrows(work: "test")
+    }
+    print("replySyncThrows (remote): \(replySyncThrows)")
+    // CHECK: replySyncThrows (remote): nil
+
+    let replyAsync = await remoteW.whenLocal { __secretlyKnownToBeLocal in
+      await __secretlyKnownToBeLocal.async(work: "test")
+    }
+    print("replyAsync (remote): \(replyAsync)")
+    // CHECK: replyAsync (remote): nil
+
+    let replyAsyncThrows = try await remoteW.whenLocal { __secretlyKnownToBeLocal in
+      try await __secretlyKnownToBeLocal.asyncThrows(work: "test")
+    }
+    print("replyAsyncThrows (remote): \(replyAsyncThrows)")
+    // CHECK: replyAsyncThrows (remote): nil
+  }
+  // ==== ----------------------------------------------------------------------
+
+  do {
+    let replyDistSubmit = try await localW.whenLocal { __secretlyKnownToBeLocal in
+      try await __secretlyKnownToBeLocal.submit(work: "local-test")
+    }
+    print("replyDistSubmit (local): \(replyDistSubmit ?? "nil")")
+    // CHECK-NOT: >> remoteCall
+    // CHECK: replyDistSubmit (local): submit(work:): local-test
+
+    let replySyncLocal = await localW.whenLocal { __secretlyKnownToBeLocal in
+      __secretlyKnownToBeLocal.sync(work: "local-test")
+    }
+    print("replySyncLocal (local): \(replySyncLocal ?? "nil")")
+    // CHECK-NOT: >> remoteCall
+    // CHECK: replySyncLocal (local): sync(work:): local-test
+
+    let replySyncThrowsLocal = try await localW.whenLocal { __secretlyKnownToBeLocal in
+      try __secretlyKnownToBeLocal.syncThrows(work: "local-test")
+    }
+    print("replySyncThrowsLocal (local): \(replySyncThrowsLocal ?? "nil")")
+    // CHECK-NOT: >> remoteCall
+    // CHECK: replySyncThrowsLocal (local): syncThrows(work:): local-test
+
+    let replyAsyncLocal = await localW.whenLocal { __secretlyKnownToBeLocal in
+      await __secretlyKnownToBeLocal.async(work: "local-test")
+    }
+    print("replyAsyncLocal (local): \(replyAsyncLocal ?? "nil")")
+    // CHECK-NOT: >> remoteCall
+    // CHECK: replyAsyncLocal (local): async(work:): local-test
+
+    let replyAsyncThrowsLocal = try await localW.whenLocal { __secretlyKnownToBeLocal in
+      try await __secretlyKnownToBeLocal.asyncThrows(work: "local-test")
+    }
+    print("replyAsyncThrowsLocal (local): \(replyAsyncThrowsLocal ?? "nil")")
+    // CHECK-NOT: >> remoteCall
+    // CHECK: replyAsyncThrowsLocal (local): asyncThrows(work:): local-test
+  }
+}
+
+@main struct Main {
+  static func main() async {
+    let system = DefaultDistributedActorSystem()
+    try! await test_generic(system: system)
+    print("==== ---------------------------------------------------")
+    try! await test_whenLocal(system: system)
+  }
+}