AssociatedTypeInference: Allow inference in compatible constrained extensions

lib/Sema/TypeCheckProtocolInference.cpp

@@ -171,35 +171,39 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
 
   InferredAssociatedTypesByWitnesses result;
 
-  auto isExtensionUsableForInference = [&](ExtensionDecl *extension) -> bool {
-
-    // The extension where the conformance being checked is declared.
-    auto conformanceExtension = checker.Conformance->
-      getDeclContext()->getAsDecl();
-    if (extension == conformanceExtension)
+  auto isExtensionUsableForInference = [&](const ExtensionDecl *extension) {
+    // The context the conformance being checked is declared on.
+    const auto conformanceCtx = checker.Conformance->getDeclContext();
+    if (extension == conformanceCtx)
       return true;
 
-    auto *extendedNominal = extension->getExtendedNominal();
-
     // Invalid case.
+    const auto extendedNominal = extension->getExtendedNominal();
     if (extendedNominal == nullptr)
       return true;
-
-    // Assume unconstrained concrete extensions we found witnesses in are
-    // always viable.
-    if (!isa<ProtocolDecl>(extendedNominal))
-      return !extension->isConstrainedExtension();
-
+
     // FIXME: The extension may not have a generic signature set up yet as
     // resolving signatures may trigger associated type inference.  This cycle
     // is now detectable and we should look into untangling it
     // - see rdar://55263708
     if (!extension->hasComputedGenericSignature())
       return true;
 
-    // Build a generic signature.
-    auto extensionSig = extension->getGenericSignature();
-
+    // Retrieve the generic signature of the extension.
+    const auto extensionSig = extension->getGenericSignature();
+
+    // If the extension is bound to the nominal the conformance is
+    // declared on, it is viable for inference when its conditional
+    // requirements are satisfied by those of the conformance context.
+    if (!isa<ProtocolDecl>(extendedNominal)) {
+      // Extensions of non-generic nominals are always viable for inference.
+      if (!extensionSig)
+        return true;
+
+      return extensionSig->requirementsNotSatisfiedBy(
+          conformanceCtx->getGenericSignatureOfContext()).empty();
+    }
+
     // The condition here is a bit more fickle than
     // `isExtensionApplied`. That check would prematurely reject
     // extensions like `P where AssocType == T` if we're relying on a
@@ -413,9 +417,7 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
   InferredAssociatedTypes result;
   for (auto member : proto->getMembers()) {
     auto req = dyn_cast<ValueDecl>(member);
-    if (!req)
-      continue;
-    if (!req->isProtocolRequirement())
+    if (!req || !req->isProtocolRequirement())
       continue;
 
     // Infer type witnesses for associated types.
@@ -449,19 +451,14 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
 
     // Check whether any of the associated types we care about are
     // referenced in this value requirement.
-    bool anyAssocTypeMatches = false;
-    for (auto assocType : checker.getReferencedAssociatedTypes(req)) {
-      if (assocTypes.count(assocType) > 0) {
-        anyAssocTypeMatches = true;
-        break;
-      }
+    {
+      const auto referenced = checker.getReferencedAssociatedTypes(req);
+      if (llvm::find_if(referenced, [&](AssociatedTypeDecl *const assocType) {
+                          return assocTypes.count(assocType);
+                        }) == referenced.end())
+        continue;
     }
 
-    // We cannot deduce anything from the witnesses of this
-    // requirement; skip it.
-    if (!anyAssocTypeMatches)
-      continue;
-
     // Infer associated types from the potential value witnesses for
     // this requirement.
     auto reqInferred =

test/decl/protocol/req/associated_type_inference.swift

@@ -527,3 +527,38 @@ extension S30 {
     T.bar()
   }
 }
+
+protocol P32 {
+  associatedtype A
+  associatedtype B
+  associatedtype C
+
+  func foo(arg: A) -> C
+  var bar: B { get }
+}
+protocol P33 {
+  associatedtype A
+
+  var baz: A { get } // expected-note {{protocol requires property 'baz' with type 'S31<T>.A' (aka 'Never'); do you want to add a stub?}}
+}
+protocol P34 {
+  associatedtype A
+
+  func boo() -> A // expected-note {{protocol requires function 'boo()' with type '() -> S31<T>.A' (aka '() -> Never'); do you want to add a stub?}}
+}
+struct S31<T> {}
+extension S31: P32 where T == Int {} // OK
+extension S31 where T == Int {
+  func foo(arg: Never) {}
+}
+extension S31 where T: Equatable {
+  var bar: Bool { true }
+}
+extension S31: P33 where T == Never {} // expected-error {{type 'S31<T>' does not conform to protocol 'P33'}}
+extension S31 where T == String {
+  var baz: Bool { true } // expected-note {{candidate has non-matching type 'Bool' [with A = S31<T>.A]}}
+}
+extension S31: P34 {} // expected-error {{type 'S31<T>' does not conform to protocol 'P34'}}
+extension S31 where T: P32 {
+  func boo() -> Void {} // expected-note {{candidate has non-matching type '<T> () -> Void' [with A = S31<T>.A]}}
+}