handle conformance requirement on extension in distributed funcs

Author: ktoso

include/swift/AST/DistributedDecl.h

@@ -50,7 +50,8 @@ 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);
+Type getSerializationRequirementTypesForMember(
+    ValueDecl *member, llvm::SmallPtrSet<ProtocolDecl *, 2> &serializationRequirements);
 
 /// Get specific 'SerializationRequirement' as defined in 'nominal'
 /// type, which must conform to the passed 'protocol' which is expected
@@ -114,17 +115,6 @@ getDistributedSerializationRequirements(
     ProtocolDecl *protocol,
     llvm::SmallPtrSetImpl<ProtocolDecl *> &requirementProtos);
 
-/// Given any set of generic requirements, locate those which are about the
-/// `SerializationRequirement`. Those need to be applied in the parameter and
-/// return type checking of distributed targets.
-void
-extractDistributedSerializationRequirements(
-    ASTContext &C,
-    ArrayRef<Requirement> allRequirements,
-    llvm::SmallPtrSet<ProtocolDecl *, 2> &into);
-
-}
-
 // ==== ------------------------------------------------------------------------
 
 #endif /* SWIFT_DECL_DISTRIBUTEDDECL_H */

lib/AST/DistributedDecl.cpp

@@ -95,8 +95,9 @@ Type swift::getConcreteReplacementForProtocolActorSystemType(ValueDecl *member)
   llvm_unreachable("Unable to fetch ActorSystem type!");
 }
 
-Type swift::getConcreteReplacementForMemberSerializationRequirement(
-    ValueDecl *member) {
+Type swift::getSerializationRequirementTypesForMember(
+    ValueDecl *member,
+    llvm::SmallPtrSet<ProtocolDecl *, 2> &serializationRequirements) {
   auto &C = member->getASTContext();
   auto *DC = member->getDeclContext();
   auto DA = C.getDistributedActorDecl();
@@ -117,6 +118,18 @@ Type swift::getConcreteReplacementForMemberSerializationRequirement(
       signature = DC->getGenericSignatureOfContext();
     }
 
+    // Also store all `SerializationRequirement : SomeProtocol` requirements
+    for (auto requirement: signature.getRequirements()) {
+      if (requirement.getFirstType()->isEqual(SerReqAssocType) &&
+          requirement.getKind() == RequirementKind::Conformance) {
+        if (auto nominal = requirement.getSecondType()->getAnyNominal()) {
+          if (auto protocol = dyn_cast<ProtocolDecl>(nominal)) {
+            serializationRequirements.insert(protocol);
+          }
+        }
+      }
+    }
+
     // 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);
@@ -1222,33 +1235,6 @@ AbstractFunctionDecl::isDistributedTargetInvocationResultHandlerOnReturn() const
     return true;
 }
 
-void
-swift::extractDistributedSerializationRequirements(
-    ASTContext &C,
-    ArrayRef<Requirement> allRequirements,
-    llvm::SmallPtrSet<ProtocolDecl *, 2> &into) {
-  auto DA = C.getDistributedActorDecl();
-  auto daSerializationReqAssocType =
-      DA->getAssociatedType(C.Id_SerializationRequirement);
-
-  for (auto req : allRequirements) {
-    // FIXME: Seems unprincipled
-    if (req.getKind() != RequirementKind::SameType &&
-        req.getKind() != RequirementKind::Conformance)
-      continue;
-
-    if (auto dependentMemberType =
-            req.getFirstType()->getAs<DependentMemberType>()) {
-      if (dependentMemberType->getAssocType() == daSerializationReqAssocType) {
-        auto layout = req.getSecondType()->getExistentialLayout();
-        for (auto p : layout.getProtocols()) {
-          serializationReqs.insert(p);
-        }
-      }
-    }
-  }
-}
-
 /******************************************************************************/
 /********************** Distributed Functions *********************************/
 /******************************************************************************/

lib/Sema/TypeCheckDistributed.cpp

@@ -386,11 +386,16 @@ bool swift::checkDistributedActorSystemAdHocProtocolRequirements(
 static bool checkDistributedTargetResultType(
     ModuleDecl *module, ValueDecl *valueDecl,
     Type serializationRequirement,
+    llvm::SmallPtrSet<ProtocolDecl *, 2> serializationRequirements,
     bool diagnose) {
   auto &C = valueDecl->getASTContext();
 
-  if (!serializationRequirement || serializationRequirement->hasError())
+  if (serializationRequirement && serializationRequirement->hasError()) {
+    return false;
+  }
+  if ((!serializationRequirement || serializationRequirement->hasError()) && serializationRequirements.empty()) {
     return false; // error of the type would be diagnosed elsewhere
+  }
 
   Type resultType;
   if (auto func = dyn_cast<FuncDecl>(valueDecl)) {
@@ -404,37 +409,43 @@ static bool checkDistributedTargetResultType(
   if (resultType->isVoid())
     return false;
 
+
+  // Collect extra "SerializationRequirement: SomeProtocol" requirements
+  if (serializationRequirement && !serializationRequirement->hasError()) {
+    auto srl = serializationRequirement->getExistentialLayout();
+    for (auto s: srl.getProtocols()) {
+      serializationRequirements.insert(s);
+    }
+  }
+
   auto isCodableRequirement =
       checkDistributedSerializationRequirementIsExactlyCodable(
           C, serializationRequirement);
 
-  if (serializationRequirement && !serializationRequirement->hasError()) {
-    auto srl = serializationRequirement->getExistentialLayout();
-    for (auto serializationReq: srl.getProtocols()) {
-      auto conformance =
-          TypeChecker::conformsToProtocol(resultType, serializationReq, module);
-      if (conformance.isInvalid()) {
-        if (diagnose) {
-          llvm::StringRef conformanceToSuggest = isCodableRequirement ?
-                                                 "Codable" : // Codable is a typealias, easier to diagnose like that
-                                                 serializationReq->getNameStr();
-
-          auto diag = valueDecl->diagnose(
-              diag::distributed_actor_target_result_not_codable,
-              resultType,
-              valueDecl,
-              conformanceToSuggest
-          );
-
-          if (isCodableRequirement) {
-            if (auto resultNominalType = resultType->getAnyNominal()) {
-              addCodableFixIt(resultNominalType, diag);
-            }
+  for (auto serializationReq: serializationRequirements) {
+    auto conformance =
+        TypeChecker::conformsToProtocol(resultType, serializationReq, module);
+    if (conformance.isInvalid()) {
+      if (diagnose) {
+        llvm::StringRef conformanceToSuggest = isCodableRequirement ?
+                                               "Codable" : // Codable is a typealias, easier to diagnose like that
+                                               serializationReq->getNameStr();
+
+        auto diag = valueDecl->diagnose(
+            diag::distributed_actor_target_result_not_codable,
+            resultType,
+            valueDecl,
+            conformanceToSuggest
+        );
+
+        if (isCodableRequirement) {
+          if (auto resultNominalType = resultType->getAnyNominal()) {
+            addCodableFixIt(resultNominalType, diag);
           }
-        } // end if: diagnose
+        }
+      } // end if: diagnose
 
-        return true;
-      }
+      return true;
     }
   }
 
@@ -502,16 +513,16 @@ bool CheckDistributedFunctionRequest::evaluate(
   }
 
   auto &C = func->getASTContext();
-  auto DC = func->getDeclContext();
   auto module = func->getParentModule();
 
   /// If no distributed module is available, then no reason to even try checks.
   if (!C.getLoadedModule(C.Id_Distributed))
     return true;
 
-  Type serializationReqType = getConcreteReplacementForMemberSerializationRequirement(func);
-  for (auto param: *func->getParameters()) {
+  llvm::SmallPtrSet<ProtocolDecl *, 2> serializationRequirements;
+  Type serializationReqType = getSerializationRequirementTypesForMember(func, serializationRequirements);
 
+  for (auto param: *func->getParameters()) {
     // --- Check the parameter conforming to serialization requirements
     if (serializationReqType && !serializationReqType->hasError()) {
       // If the requirement is exactly `Codable` we diagnose it ia bit nicer.
@@ -574,12 +585,11 @@ bool CheckDistributedFunctionRequest::evaluate(
     }
   }
 
-  if (serializationReqType && !serializationReqType->hasError()) {
-    // --- Result type must be either void or a codable type
-    if (checkDistributedTargetResultType(module, func, serializationReqType,
-        /*diagnose=*/true)) {
-      return true;
-    }
+  // --- Result type must be either void or a serialization requirement conforming type
+  if (checkDistributedTargetResultType(
+      module, func, serializationReqType, serializationRequirements,
+      /*diagnose=*/true)) {
+    return true;
   }
 
   return false;
@@ -627,15 +637,16 @@ bool swift::checkDistributedActorProperty(VarDecl *var, bool diagnose) {
       DC->getSelfNominalTypeDecl()->getDistributedActorSystemProperty();
   auto systemDecl = systemVar->getInterfaceType()->getAnyNominal();
 
-//  auto serializationRequirements =
-//      getDistributedSerializationRequirementProtocols(
-//          systemDecl,
-//          C.getProtocol(KnownProtocolKind::DistributedActorSystem));
+  auto serializationRequirements =
+      getDistributedSerializationRequirementProtocols(
+          systemDecl,
+          C.getProtocol(KnownProtocolKind::DistributedActorSystem));
+
   auto serializationRequirement =
-      getConcreteReplacementForMemberSerializationRequirement(systemVar);
+      getSerializationRequirementTypesForMember(systemVar, serializationRequirements);
 
   auto module = var->getModuleContext();
-  if (checkDistributedTargetResultType(module, var, serializationRequirement, diagnose)) {
+  if (checkDistributedTargetResultType(module, var, serializationRequirement, serializationRequirements, diagnose)) {
     return true;
   }
 

test/Distributed/distributed_protocols_distributed_func_serialization_requirements.swift

@@ -82,7 +82,7 @@ extension NoSerializationRequirementYet
 
 extension NoSerializationRequirementYet
   where SerializationRequirement: Codable {
-  // expected-error@+1{{result type 'NotCodable' of distributed instance method 'test4' does not conform to serialization requirement 'Codable'}}
+  // expected-error@+1{{result type 'NotCodable' of distributed instance method 'test4' does not conform to serialization requirement 'Decodable'}}
   distributed func test4() -> NotCodable {
     .init()
   }