Sema: Generalize ProtocolDecl::getPrimaryAssociatedType() to ProtocolDecl::getPrimaryAssociatedTypes()

Author: slavapestov

include/swift/AST/Decl.h

@@ -4402,10 +4402,10 @@ class ProtocolDecl final : public NominalTypeDecl {
   /// a protocol having nested types (ObjC protocols).
   ArrayRef<AssociatedTypeDecl *> getAssociatedTypeMembers() const;
 
-  /// Returns the primary associated type, or nullptr if there isn't one. This is
-  /// the associated type that is parametrized with a same-type requirement in a
-  /// parametrized protocol type of the form SomeProtocol<SomeArgType>.
-  AssociatedTypeDecl *getPrimaryAssociatedType() const;
+  /// Returns the list of primary associated types. These are the associated
+  /// types that is parametrized with same-type requirements in a
+  /// parametrized protocol type of the form SomeProtocol<Arg1, Arg2...>.
+  ArrayRef<AssociatedTypeDecl *> getPrimaryAssociatedTypes() const;
 
   /// Returns a protocol requirement with the given name, or nullptr if the
   /// name has multiple overloads, or no overloads at all.

include/swift/AST/DiagnosticsSema.def

@@ -3734,11 +3734,13 @@ ERROR(not_a_generic_definition,none,
 ERROR(not_a_generic_type,none,
       "cannot specialize non-generic type %0", (Type))
 ERROR(parameterized_protocol_not_supported,none,
-      "protocol type with generic argument can only be used as a generic constraint", ())
+      "protocol type with generic arguments can only be used as a generic constraint", ())
 ERROR(protocol_does_not_have_primary_assoc_type,none,
       "cannot specialize protocol type %0", (Type))
-ERROR(protocol_cannot_have_multiple_generic_arguments,none,
-      "protocol type %0 can only be specialized with exactly one argument", (Type))
+ERROR(parameterized_protocol_too_many_type_arguments,none,
+      "protocol type %0 can specialized with too many type parameters "
+      "(got %1, but expected at most %2)",
+      (Type, unsigned, unsigned))
 ERROR(cannot_specialize_self,none,
       "cannot specialize 'Self'", ())
 NOTE(specialize_explicit_type_instead,none,

include/swift/AST/TypeCheckRequests.h

@@ -313,10 +313,10 @@ class ExistentialRequiresAnyRequest :
   void cacheResult(bool value) const;
 };
 
-/// Find the primary associated type of the given protocol.
-class PrimaryAssociatedTypeRequest :
-    public SimpleRequest<PrimaryAssociatedTypeRequest,
-                         AssociatedTypeDecl *(ProtocolDecl *),
+/// Find the list of primary associated types of the given protocol.
+class PrimaryAssociatedTypesRequest :
+    public SimpleRequest<PrimaryAssociatedTypesRequest,
+                         ArrayRef<AssociatedTypeDecl *>(ProtocolDecl *),
                          RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -325,7 +325,7 @@ class PrimaryAssociatedTypeRequest :
   friend SimpleRequest;
 
   // Evaluation.
-  AssociatedTypeDecl *
+  ArrayRef<AssociatedTypeDecl *>
   evaluate(Evaluator &evaluator, ProtocolDecl *decl) const;
 
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -262,8 +262,8 @@ SWIFT_REQUEST(TypeChecker, PropertyWrapperTypeInfoRequest,
               NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ProtocolRequiresClassRequest, bool(ProtocolDecl *),
               SeparatelyCached, NoLocationInfo)
-SWIFT_REQUEST(TypeChecker, PrimaryAssociatedTypeRequest,
-              AssociatedTypeDecl *(ProtocolDecl *),
+SWIFT_REQUEST(TypeChecker, PrimaryAssociatedTypesRequest,
+              ArrayRef<AssociatedTypeDecl *>(ProtocolDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, RequirementRequest,
               Requirement(WhereClauseOwner, unsigned, TypeResolutionStage),

include/swift/AST/Types.h

@@ -5289,6 +5289,8 @@ class ParameterizedProtocolType final : public TypeBase,
             Bits.ParameterizedProtocolType.ArgCount};
   }
 
+  void getRequirements(Type baseType, SmallVectorImpl<Requirement> &reqs) const;
+
   void Profile(llvm::FoldingSetNodeID &ID) {
     Profile(ID, Base, getArgs());
   }

lib/AST/Decl.cpp

@@ -5398,9 +5398,10 @@ bool ProtocolDecl::existentialRequiresAny() const {
     ExistentialRequiresAnyRequest{const_cast<ProtocolDecl *>(this)}, true);
 }
 
-AssociatedTypeDecl *ProtocolDecl::getPrimaryAssociatedType() const {
+ArrayRef<AssociatedTypeDecl *>
+ProtocolDecl::getPrimaryAssociatedTypes() const {
   return evaluateOrDefault(getASTContext().evaluator,
-    PrimaryAssociatedTypeRequest{const_cast<ProtocolDecl *>(this)},
+    PrimaryAssociatedTypesRequest{const_cast<ProtocolDecl *>(this)},
     nullptr);
 }
 

lib/AST/GenericSignatureBuilder.cpp

@@ -4523,14 +4523,13 @@ ConstraintResult GenericSignatureBuilder::addTypeRequirement(
                                                 source)))
         anyErrors = true;
 
-    auto *assocType = protoDecl->getPrimaryAssociatedType();
-    auto depType = DependentMemberType::get(
-        resolvedSubject.getDependentType(*this), assocType);
-    if (isErrorResult(addSameTypeRequirement(Type(depType),
-                                             paramProtoType->getArgs()[0],
-                                             source,
-                                             UnresolvedHandlingKind::GenerateConstraints)))
-      anyErrors = true;
+    SmallVector<Requirement, 2> reqs;
+    auto baseType = resolvedSubject.getDependentType(*this);
+    paramProtoType->getRequirements(baseType, reqs);
+    for (auto req : reqs) {
+      if (isErrorResult(addRequirement(req, source, inferForModule)))
+        anyErrors = true;
+    }
 
     return anyErrors ? ConstraintResult::Conflicting
                      : ConstraintResult::Resolved;

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -217,16 +217,15 @@ static void desugarConformanceRequirement(Type subjectType, Type constraintType,
   }
 
   if (auto *paramType = constraintType->getAs<ParameterizedProtocolType>()) {
-    auto *protoDecl = paramType->getBaseType()->getDecl();
-
     desugarConformanceRequirement(subjectType, paramType->getBaseType(),
                                   loc, result, errors);
 
-    auto *assocType = protoDecl->getPrimaryAssociatedType();
+    SmallVector<Requirement, 2> reqs;
+    paramType->getRequirements(subjectType, reqs);
+
+    for (const auto &req : reqs)
+      desugarRequirement(req, result, errors);
 
-    auto memberType = lookupMemberType(subjectType, protoDecl, assocType);
-    desugarSameTypeRequirement(memberType, paramType->getArgs()[0],
-                               loc, result, errors);
     return;
   }
 

lib/AST/Type.cpp

@@ -3894,6 +3894,24 @@ void ParameterizedProtocolType::Profile(llvm::FoldingSetNodeID &ID,
     ID.AddPointer(arg.getPointer());
 }
 
+void ParameterizedProtocolType::getRequirements(
+    Type baseType, SmallVectorImpl<Requirement> &reqs) const {
+  auto *protoDecl = getBaseType()->getDecl();
+
+  auto assocTypes = protoDecl->getPrimaryAssociatedTypes();
+  auto argTypes = getArgs();
+  assert(argTypes.size() <= assocTypes.size());
+
+  for (unsigned i : indices(argTypes)) {
+    auto argType = argTypes[i];
+    auto *assocType = assocTypes[i];
+    auto subjectType = assocType->getDeclaredInterfaceType()
+        ->castTo<DependentMemberType>()
+        ->substBaseType(protoDecl->getParentModule(), baseType);
+    reqs.emplace_back(RequirementKind::SameType, subjectType, argType);
+  }
+}
+
 bool ProtocolType::requiresClass() {
   return getDecl()->requiresClass();
 }

lib/Sema/TypeCheckDecl.cpp

@@ -698,15 +698,17 @@ ExistentialRequiresAnyRequest::evaluate(Evaluator &evaluator,
   return false;
 }
 
-AssociatedTypeDecl *
-PrimaryAssociatedTypeRequest::evaluate(Evaluator &evaluator,
-                                       ProtocolDecl *decl) const {
+ArrayRef<AssociatedTypeDecl *>
+PrimaryAssociatedTypesRequest::evaluate(Evaluator &evaluator,
+                                        ProtocolDecl *decl) const {
+  SmallVector<AssociatedTypeDecl *, 2> assocTypes;
+
   for (auto *assocType : decl->getAssociatedTypeMembers()) {
     if (assocType->getAttrs().hasAttribute<PrimaryAssociatedTypeAttr>())
-      return assocType;
+      assocTypes.push_back(assocType);
   }
 
-  return nullptr;
+  return decl->getASTContext().AllocateCopy(assocTypes);
 }
 
 bool