AST: Introduce ProtocolDecl::getPrimaryAssociatedType()

Author: slavapestov

include/swift/AST/Decl.h

@@ -4372,6 +4372,11 @@ 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 a protocol requirement with the given name, or nullptr if the
   /// name has multiple overloads, or no overloads at all.
   ValueDecl *getSingleRequirement(DeclName name) const;

include/swift/AST/TypeCheckRequests.h

@@ -313,6 +313,26 @@ class ExistentialRequiresAnyRequest :
   void cacheResult(bool value) const;
 };
 
+/// Find the primary associated type of the given protocol.
+class PrimaryAssociatedTypeRequest :
+    public SimpleRequest<PrimaryAssociatedTypeRequest,
+                         AssociatedTypeDecl *(ProtocolDecl *),
+                         RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  // Evaluation.
+  AssociatedTypeDecl *
+  evaluate(Evaluator &evaluator, ProtocolDecl *decl) const;
+
+public:
+  // Caching.
+  bool isCached() const { return true; }
+};
+
 class PolymorphicEffectRequirementsRequest :
     public SimpleRequest<PolymorphicEffectRequirementsRequest,
                          PolymorphicEffectRequirementList(EffectKind, ProtocolDecl *),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -233,6 +233,9 @@ SWIFT_REQUEST(TypeChecker, PropertyWrapperTypeInfoRequest,
               NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ProtocolRequiresClassRequest, bool(ProtocolDecl *),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PrimaryAssociatedTypeRequest,
+              AssociatedTypeDecl *(ProtocolDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, RequirementRequest,
               Requirement(WhereClauseOwner, unsigned, TypeResolutionStage),
               Cached, HasNearestLocation)

lib/AST/Decl.cpp

@@ -5284,6 +5284,12 @@ bool ProtocolDecl::existentialRequiresAny() const {
     ExistentialRequiresAnyRequest{const_cast<ProtocolDecl *>(this)}, true);
 }
 
+AssociatedTypeDecl *ProtocolDecl::getPrimaryAssociatedType() const {
+  return evaluateOrDefault(getASTContext().evaluator,
+    PrimaryAssociatedTypeRequest{const_cast<ProtocolDecl *>(this)},
+    nullptr);
+}
+
 StringRef ProtocolDecl::getObjCRuntimeName(
                           llvm::SmallVectorImpl<char> &buffer) const {
   // If there is an 'objc' attribute with a name, use that name.

lib/Sema/TypeCheckDecl.cpp

@@ -694,6 +694,17 @@ ExistentialRequiresAnyRequest::evaluate(Evaluator &evaluator,
   return false;
 }
 
+AssociatedTypeDecl *
+PrimaryAssociatedTypeRequest::evaluate(Evaluator &evaluator,
+                                       ProtocolDecl *decl) const {
+  for (auto *assocType : decl->getAssociatedTypeMembers()) {
+    if (assocType->getAttrs().hasAttribute<PrimaryAssociatedTypeAttr>())
+      return assocType;
+  }
+
+  return nullptr;
+}
+
 bool
 IsFinalRequest::evaluate(Evaluator &evaluator, ValueDecl *decl) const {
   if (isa<ClassDecl>(decl))