[Distributed] distributed func checks now work in protocols & dont crash

Author: ktoso

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();
@@ -788,7 +793,6 @@ FuncDecl *GetDistributedThunkRequest::evaluate(
     return nullptr;
 
   auto &C = distributedTarget->getASTContext();
-  auto DC = distributedTarget->getDeclContext();
 
   if (!getConcreteReplacementForProtocolActorSystemType(distributedTarget)) {
     // Don't synthesize thunks, unless there is a *concrete* ActorSystem.
@@ -812,9 +816,6 @@ FuncDecl *GetDistributedThunkRequest::evaluate(
     if (!C.getLoadedModule(C.Id_Distributed))
       return nullptr;
 
-    auto nominal = DC->getSelfNominalTypeDecl(); // NOTE: Always from DC
-    assert(nominal);
-
     // --- Prepare the "distributed thunk" which does the "maybe remote" dance:
     return createDistributedThunkFunction(func);
   }

lib/Sema/TypeCheckAttr.cpp

@@ -5900,16 +5900,32 @@ void AttributeChecker::visitDistributedActorAttr(DistributedActorAttr *attr) {
 
     // distributed func must be declared inside an distributed actor
     auto selfTy = dc->getSelfTypeInContext();
+
+    auto *protoDecl = dc->getSelfProtocolDecl();
     if (!selfTy->isDistributedActor()) {
       auto diagnostic = diagnoseAndRemoveAttr(
         attr, diag::distributed_actor_func_not_in_distributed_actor);
 
-      if (auto *protoDecl = dc->getSelfProtocolDecl()) {
+      if (protoDecl) {
         diagnoseDistributedFunctionInNonDistributedActorProtocol(protoDecl,
                                                                  diagnostic);
       }
       return;
     }
+
+    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/TypeCheckDeclPrimary.cpp

@@ -23,6 +23,7 @@
 #include "TypeCheckAvailability.h"
 #include "TypeCheckConcurrency.h"
 #include "TypeCheckDecl.h"
+#include "TypeCheckDistributed.h"
 #include "TypeCheckObjC.h"
 #include "TypeCheckType.h"
 #include "TypeChecker.h"
@@ -33,6 +34,7 @@
 #include "swift/AST/AccessScope.h"
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/DistributedDecl.h"
 #include "swift/AST/ForeignErrorConvention.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Initializer.h"
@@ -2875,6 +2877,7 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
     }
 
     TypeChecker::checkDeclAttributes(FD);
+    TypeChecker::checkDistributedFunc(FD);
 
     if (!checkOverrides(FD)) {
       // If a method has an 'override' keyword but does not

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);
       }
@@ -676,6 +703,13 @@ void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal
   (void)nominal->getDistributedActorIDProperty();
 }
 
+void TypeChecker::checkDistributedFunc(FuncDecl *func) {
+  if (!func->isDistributed())
+    return;
+
+  swift::checkDistributedFunction(func);
+}
+
 llvm::SmallPtrSet<ProtocolDecl *, 2>
 swift::getDistributedSerializationRequirementProtocols(
     NominalTypeDecl *nominal, ProtocolDecl *protocol) {

lib/Sema/TypeCheckProtocol.cpp

@@ -1401,9 +1401,10 @@ bool WitnessChecker::findBestWitness(
        attempt = static_cast<Attempt>(attempt + 1)) {
     SmallVector<ValueDecl *, 4> witnesses;
     switch (attempt) {
-    case Regular:
+    case Regular: {
       witnesses = lookupValueWitnesses(requirement, ignoringNames);
       break;
+    }
     case OperatorsFromOverlay: {
       // If we have a Clang declaration, the matching operator might be in the
       // overlay for that module.

lib/Sema/TypeChecker.h

@@ -1101,6 +1101,9 @@ diagnosePotentialOpaqueTypeUnavailability(SourceRange ReferenceRange,
 /// Type check a 'distributed actor' declaration.
 void checkDistributedActor(SourceFile *SF, NominalTypeDecl *decl);
 
+/// Type check a single 'distributed func' declaration.
+void checkDistributedFunc(FuncDecl *func);
+
 void checkConcurrencyAvailability(SourceRange ReferenceRange,
                                   const DeclContext *ReferenceDC);