Use more of getConcreteReplacementForMemberSerializationRequirement

Author: ktoso

include/swift/AST/DistributedDecl.h

@@ -97,7 +97,7 @@ getDistributedSerializationRequirementProtocols(
 /// If so, we can emit slightly nicer diagnostics.
 bool checkDistributedSerializationRequirementIsExactlyCodable(
     ASTContext &C,
-    const llvm::SmallPtrSetImpl<ProtocolDecl *> &allRequirements);
+    Type type);
 
 /// Get the `SerializationRequirement`, explode it into the specific
 /// protocol requirements and insert them into `requirements`.

lib/AST/DistributedDecl.cpp

@@ -106,18 +106,17 @@ Type swift::getConcreteReplacementForMemberSerializationRequirement(
     return getDistributedSerializationRequirementType(classDecl, C.getDistributedActorDecl());
   }
 
-  /// === Maybe the value is declared in a protocol?
-  if (auto protocol = DC->getSelfProtocolDecl()) {
+  auto SerReqAssocType = DA->getAssociatedType(C.Id_SerializationRequirement)
+      ->getDeclaredInterfaceType();
+
+  if (DC->getSelfProtocolDecl() || isa<ExtensionDecl>(DC)) {
     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);
@@ -355,15 +354,24 @@ swift::getDistributedSerializationRequirements(
 
 bool swift::checkDistributedSerializationRequirementIsExactlyCodable(
     ASTContext &C,
-    const llvm::SmallPtrSetImpl<ProtocolDecl *> &allRequirements) {
+    Type type) {
+  if (!type)
+    return false;
+
+  if (type->hasError())
+    return false;
+
   auto encodable = C.getProtocol(KnownProtocolKind::Encodable);
   auto decodable = C.getProtocol(KnownProtocolKind::Decodable);
 
-  if (allRequirements.size() != 2)
+  auto layout = type->getExistentialLayout();
+  auto protocols = layout.getProtocols();
+
+  if (protocols.size() != 2)
     return false;
 
-  return allRequirements.count(encodable) &&
-         allRequirements.count(decodable);
+  return std::count(protocols.begin(), protocols.end(), encodable) == 1 &&
+      std::count(protocols.begin(), protocols.end(), decodable) == 1;
 }
 
 /******************************************************************************/

lib/Sema/TypeCheckDistributed.cpp

@@ -385,10 +385,13 @@ bool swift::checkDistributedActorSystemAdHocProtocolRequirements(
 
 static bool checkDistributedTargetResultType(
     ModuleDecl *module, ValueDecl *valueDecl,
-    const llvm::SmallPtrSetImpl<ProtocolDecl *> &serializationRequirements,
+    Type serializationRequirement,
     bool diagnose) {
   auto &C = valueDecl->getASTContext();
 
+  if (!serializationRequirement || serializationRequirement->hasError())
+    return false; // error of the type would be diagnosed elsewhere
+
   Type resultType;
   if (auto func = dyn_cast<FuncDecl>(valueDecl)) {
     resultType = func->mapTypeIntoContext(func->getResultInterfaceType());
@@ -403,36 +406,39 @@ static bool checkDistributedTargetResultType(
 
   auto isCodableRequirement =
       checkDistributedSerializationRequirementIsExactlyCodable(
-          C, serializationRequirements);
-
-  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);
+          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);
+            }
           }
-        }
-      } // end if: diagnose
-      
-      return true;
+        } // end if: diagnose
+
+        return true;
+      }
     }
   }
 
-  return false;
+    return false;
 }
 
 bool swift::checkDistributedActorSystem(const NominalTypeDecl *system) {
@@ -503,74 +509,35 @@ bool CheckDistributedFunctionRequest::evaluate(
   if (!C.getLoadedModule(C.Id_Distributed))
     return true;
 
-  // === All parameters and the result type must conform
-  // SerializationRequirement
-  llvm::SmallPtrSet<ProtocolDecl *, 2> serializationRequirements;
-  if (auto extension = dyn_cast<ExtensionDecl>(DC)) {
-    auto actorOrProtocol = extension->getExtendedNominal();
-    if (auto actor = dyn_cast<ClassDecl>(actorOrProtocol)) {
-      assert(actor->isAnyActor());
-      serializationRequirements = getDistributedSerializationRequirementProtocols(
-          getDistributedActorSystemType(actor)->getAnyNominal(),
-          C.getProtocol(KnownProtocolKind::DistributedActorSystem));
-    } else if (auto protocol = dyn_cast<ProtocolDecl>(actorOrProtocol)) {
-      extractDistributedSerializationRequirements(
-          C, protocol->getGenericRequirements(),
-          /*into=*/serializationRequirements);
-      extractDistributedSerializationRequirements(
-          C, extension->getGenericRequirements(),
-          /*into=*/serializationRequirements);
-    } else {
-      // ignore
-    }
-  } else if (auto actor = dyn_cast<ClassDecl>(DC)) {
-    serializationRequirements = getDistributedSerializationRequirementProtocols(
-        getDistributedActorSystemType(actor)->getAnyNominal(),
-        C.getProtocol(KnownProtocolKind::DistributedActorSystem));
-  } else if (isa<ProtocolDecl>(DC)) {
-    if (auto seqReqTy =
-        getConcreteReplacementForMemberSerializationRequirement(func)) {
-      auto layout = seqReqTy->getExistentialLayout();
-      for (auto req : layout.getProtocols()) {
-        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("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.
-  auto serializationRequirementIsCodable =
-      checkDistributedSerializationRequirementIsExactlyCodable(
-          C, serializationRequirements);
-
-  for (auto param : *func->getParameters()) {
-    // --- Check parameters for 'Codable' conformance
-    auto paramTy = func->mapTypeIntoContext(param->getInterfaceType());
-
-    for (auto req : serializationRequirements) {
-      if (TypeChecker::conformsToProtocol(paramTy, req, module).isInvalid()) {
-        auto diag = func->diagnose(
-            diag::distributed_actor_func_param_not_codable,
-            param->getArgumentName().str(), param->getInterfaceType(),
-            func->getDescriptiveKind(),
-            serializationRequirementIsCodable ? "Codable"
-                                              : req->getNameStr());
-
-        if (auto paramNominalTy = paramTy->getAnyNominal()) {
-          addCodableFixIt(paramNominalTy, diag);
-        } // else, no nominal type to suggest the fixit for, e.g. a closure
-
-        return true;
+  Type serializationReqType = getConcreteReplacementForMemberSerializationRequirement(func);
+  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.
+      auto serializationRequirementIsCodable =
+          checkDistributedSerializationRequirementIsExactlyCodable(
+              C, serializationReqType);
+
+      // --- Check parameters for 'SerializationRequirement' conformance
+      auto paramTy = func->mapTypeIntoContext(param->getInterfaceType());
+
+      auto srl = serializationReqType->getExistentialLayout();
+      for (auto req: srl.getProtocols()) {
+        if (TypeChecker::conformsToProtocol(paramTy, req, module).isInvalid()) {
+          auto diag = func->diagnose(
+              diag::distributed_actor_func_param_not_codable,
+              param->getArgumentName().str(), param->getInterfaceType(),
+              func->getDescriptiveKind(),
+              serializationRequirementIsCodable ? "Codable"
+                                                : req->getNameStr());
+
+          if (auto paramNominalTy = paramTy->getAnyNominal()) {
+            addCodableFixIt(paramNominalTy, diag);
+          } // else, no nominal type to suggest the fixit for, e.g. a closure
+
+          return true;
+        }
       }
     }
 
@@ -607,10 +574,12 @@ bool CheckDistributedFunctionRequest::evaluate(
     }
   }
 
-  // --- Result type must be either void or a codable type
-  if (checkDistributedTargetResultType(module, func, serializationRequirements,
-                                       /*diagnose=*/true)) {
-    return true;
+  if (serializationReqType && !serializationReqType->hasError()) {
+    // --- Result type must be either void or a codable type
+    if (checkDistributedTargetResultType(module, func, serializationReqType,
+        /*diagnose=*/true)) {
+      return true;
+    }
   }
 
   return false;
@@ -658,13 +627,15 @@ 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);
 
   auto module = var->getModuleContext();
-  if (checkDistributedTargetResultType(module, var, serializationRequirements, diagnose)) {
+  if (checkDistributedTargetResultType(module, var, serializationRequirement, diagnose)) {
     return true;
   }
 
@@ -771,6 +742,7 @@ bool TypeChecker::checkDistributedFunc(FuncDecl *func) {
   return swift::checkDistributedFunction(func);
 }
 
+// TODO(distributed): Remove this entirely and rely on generic signature and getConcrete to implement checks
 llvm::SmallPtrSet<ProtocolDecl *, 2>
 swift::getDistributedSerializationRequirementProtocols(
     NominalTypeDecl *nominal, ProtocolDecl *protocol) {