[Associated type inference] Support @_implements on type witnesses

Extend the longstanding support for `@_implements` on witnesses to also
work on type witnesses, as part of associated type inference. This permits
an associated type witness to have a different name from the associated
type itself.

Author: DougGregor

lib/Sema/AssociatedTypeInference.cpp

@@ -429,7 +429,8 @@ static ResolveWitnessResult resolveTypeWitnessViaLookup(
     abort();
   }
 
-  NLOptions subOptions = (NL_QualifiedDefault | NL_OnlyTypes | NL_ProtocolMembers);
+  NLOptions subOptions = (NL_QualifiedDefault | NL_OnlyTypes |
+                          NL_ProtocolMembers | NL_IncludeAttributeImplements);
 
   // Look for a member type with the same name as the associated type.
   SmallVector<ValueDecl *, 4> candidates;
@@ -455,6 +456,16 @@ static ResolveWitnessResult resolveTypeWitnessViaLookup(
     if (isa<AssociatedTypeDecl>(typeDecl))
       continue;
 
+    // If the name doesn't match and there's no appropriate @_implements
+    // attribute, skip this candidate.
+    //
+    // Also skip candidates in protocol extensions, because they tend to cause
+    // request cycles. We'll look at those during associated type inference.
+    if (assocType->getName() != typeDecl->getName() &&
+        !(witnessHasImplementsAttrForRequiredName(typeDecl, assocType) &&
+          !typeDecl->getDeclContext()->getSelfProtocolDecl()))
+      continue;
+
     auto *genericDecl = cast<GenericTypeDecl>(typeDecl);
 
     // If the declaration has generic parameters, it cannot witness an
@@ -2039,7 +2050,8 @@ AssociatedTypeInference::inferTypeWitnessesViaAssociatedType(
 
   NLOptions subOptions = (NL_QualifiedDefault |
                           NL_OnlyTypes |
-                          NL_ProtocolMembers);
+                          NL_ProtocolMembers |
+                          NL_IncludeAttributeImplements);
 
   // Look for types with the given default name that have appropriate
   // @_implements attributes.
@@ -2061,6 +2073,12 @@ AssociatedTypeInference::inferTypeWitnessesViaAssociatedType(
     if (!defaultProto)
       continue;
 
+    // If the name doesn't match and there's no appropriate @_implements
+    // attribute, skip this candidate.
+    if (defaultName.getBaseName() != typeDecl->getName() &&
+        !witnessHasImplementsAttrForRequiredName(typeDecl, assocType))
+      continue;
+
     // Determine the witness type.
     Type witnessType = getWitnessTypeForMatching(conformance, typeDecl);
     if (!witnessType) continue;

lib/Sema/TypeCheckProtocol.cpp

@@ -1281,18 +1281,18 @@ swift::matchWitness(WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
   return matchWitness(dc, req, witness, setup, matchTypes, finalize);
 }
 
-static bool
-witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
-                                        ValueDecl *requirement) {
+bool
+swift::witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
+                                               ValueDecl *requirement) {
   if (auto A = witness->getAttrs().getAttribute<ImplementsAttr>()) {
     return A->getMemberName() == requirement->getName();
   }
   return false;
 }
 
-static bool
-witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
-                                            ValueDecl *requirement) {
+bool
+swift::witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
+                                                   ValueDecl *requirement) {
   assert(requirement->isProtocolRequirement());
   auto *PD = cast<ProtocolDecl>(requirement->getDeclContext());
   if (auto A = witness->getAttrs().getAttribute<ImplementsAttr>()) {

lib/Sema/TypeCheckProtocol.h

@@ -223,6 +223,16 @@ AssociatedTypeDecl *findDefaultedAssociatedType(
     DeclContext *dc, NominalTypeDecl *adoptee,
     AssociatedTypeDecl *assocType);
 
+/// Determine whether this witness has an `@_implements` attribute whose
+/// name matches that of the given requirement.
+bool witnessHasImplementsAttrForRequiredName(ValueDecl *witness,
+                                             ValueDecl *requirement);
+
+/// Determine whether this witness has an `@_implements` attribute whose name
+/// and protocol match that of the requirement exactly.
+bool witnessHasImplementsAttrForExactRequirement(ValueDecl *witness,
+                                                 ValueDecl *requirement);
+
 }
 
 #endif // SWIFT_SEMA_PROTOCOL_H

test/attr/attr_implements.swift

@@ -97,3 +97,13 @@ func falseWhenSpecificType(_ x: SpecificType) -> Bool { return x == x }
 
 assert(trueWhenJustEquatable(SpecificType()))
 assert(!falseWhenSpecificType(SpecificType()))
+
+// @_implements on associated types
+protocol PWithAssoc {
+  associatedtype A
+}
+
+struct XWithAssoc: PWithAssoc {
+  @_implements(PWithAssoc, A)
+  typealias __P_A = Int
+}