Merge pull request #37637 from slavapestov/associated-type-inference-proto-ext-fix

Fix associated type inference crash with constrained protocol extensions

Author: slavapestov

lib/Sema/TypeCheckProtocolInference.cpp

@@ -19,6 +19,7 @@
 
 #include "swift/AST/Decl.h"
 #include "swift/AST/GenericSignature.h"
+#include "swift/AST/NameLookupRequests.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SubstitutionMap.h"
 #include "swift/AST/TypeMatcher.h"
@@ -206,20 +207,22 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
     if (extendedNominal == nullptr)
       return true;
 
-    // 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;
-
-    // Retrieve the generic signature of the extension.
-    const auto extensionSig = extension->getGenericSignature();
+    auto *proto = dyn_cast<ProtocolDecl>(extendedNominal);
 
     // 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)) {
+    if (!proto) {
+      // 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;
+
+      // Retrieve the generic signature of the extension.
+      const auto extensionSig = extension->getGenericSignature();
+
       // Extensions of non-generic nominals are always viable for inference.
       if (!extensionSig)
         return true;
@@ -235,30 +238,29 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
     // in the first place. Only check conformances on the `Self` type,
     // because those have to be explicitly declared on the type somewhere
     // so won't be affected by whatever answer inference comes up with.
-    auto selfTy = extension->getSelfInterfaceType();
-    for (const Requirement &reqt : extensionSig->getRequirements()) {
-      switch (reqt.getKind()) {
-      case RequirementKind::Conformance:
-      case RequirementKind::Superclass:
-        // FIXME: This is the wrong check
-        if (selfTy->isEqual(reqt.getFirstType()) &&
-            !TypeChecker::isSubtypeOf(conformance->getType(),
-                                      reqt.getSecondType(), dc))
-          return false;
-        break;
+    auto *module = dc->getParentModule();
+    auto checkConformance = [&](ProtocolDecl *proto) {
+      auto otherConf = module->lookupConformance(conformance->getType(),
+                                                 proto);
+      return (otherConf && otherConf.getConditionalRequirements().empty());
+    };
 
-      case RequirementKind::Layout:
-      case RequirementKind::SameType:
-        break;
+    // First check the extended protocol itself.
+    if (!checkConformance(proto))
+      return false;
+
+    // Now check any additional bounds on 'Self' from the where clause.
+    auto bounds = getSelfBoundsFromWhereClause(extension);
+    for (auto *decl : bounds.decls) {
+      if (auto *proto = dyn_cast<ProtocolDecl>(decl)) {
+        if (!checkConformance(proto))
+          return false;
       }
     }
 
     return true;
   };
 
-  auto typeInContext =
-    conformance->getDeclContext()->mapTypeIntoContext(conformance->getType());
-
   for (auto witness :
        checker.lookupValueWitnesses(req, /*ignoringNames=*/nullptr)) {
     LLVM_DEBUG(llvm::dbgs() << "Inferring associated types from decl:\n";
@@ -314,6 +316,10 @@ AssociatedTypeInference::inferTypeWitnessesViaValueWitnesses(
           if (!associatedTypesAreSameEquivalenceClass(dmt->getAssocType(),
                                                       result.first))
             return false;
+
+          auto typeInContext =
+            conformance->getDeclContext()->mapTypeIntoContext(conformance->getType());
+
           if (!dmt->getBase()->isEqual(typeInContext))
             return false;
 

test/decl/protocol/req/associated_type_inference_proto_ext.swift

@@ -0,0 +1,30 @@
+// RUN: %target-typecheck-verify-swift
+
+struct Pair<A, B> {
+    var a: A
+    var b: B
+}
+
+extension Pair: Codable where A: Codable, B: Codable {}
+
+extension Pair: Collection where A == B {
+// expected-error@-1 {{conditional conformance of type 'Pair<A, B>' to protocol 'Collection' does not imply conformance to inherited protocol 'Sequence'}}
+// expected-note@-2 {{did you mean to explicitly state the conformance like 'extension Pair: Sequence where ...'?}}
+    typealias Element = A
+
+    var startIndex: Int { return 0 }
+    var endIndex: Int   { return 2 }
+
+    subscript(index: Int) -> Element {
+        switch index {
+        case 0: return self.a
+        case 1: return self.b
+        default: fatalError("Index out of bounds.")
+        }
+    }
+    
+    func index(after i: Int) -> Int {
+        precondition(i < endIndex, "Can't advance beyond endIndex")
+        return i + 1
+    }
+}