Sema: Try to resolve type witnesses on a more specific conformance first

Author: slavapestov

lib/Sema/AssociatedTypeInference.cpp

@@ -3997,6 +3997,36 @@ ResolveTypeWitnessesRequest::evaluate(Evaluator &evaluator,
   return evaluator::SideEffect();
 }
 
+static NormalProtocolConformance *
+getBetterConformanceForResolvingTypeWitnesses(NormalProtocolConformance *conformance,
+                                              AssociatedTypeDecl *requirement) {
+  auto *dc = conformance->getDeclContext();
+  assert(dc->getParentSourceFile() && "What are you doing?");
+  auto *proto = conformance->getProtocol();
+
+  IterableDeclContext *idc;
+  if (auto *extensionDecl = dyn_cast<ExtensionDecl>(dc))
+    idc = extensionDecl;
+  else
+    idc = cast<NominalTypeDecl>(dc);
+
+  auto otherConformances = idc->getLocalConformances(ConformanceLookupKind::NonStructural);
+  for (auto *otherConformance : otherConformances) {
+    auto *otherNormal = dyn_cast<NormalProtocolConformance>(
+        otherConformance->getRootConformance());
+    if (otherNormal == nullptr)
+      continue;
+
+    auto *otherProto = otherNormal->getProtocol();
+    if (otherProto->inheritsFrom(proto) &&
+        otherProto->getAssociatedType(requirement->getName())) {
+      return otherNormal;
+    }
+  }
+
+  return conformance;
+}
+
 TypeWitnessAndDecl
 TypeWitnessRequest::evaluate(Evaluator &eval,
                              NormalProtocolConformance *conformance,
@@ -4008,8 +4038,30 @@ TypeWitnessRequest::evaluate(Evaluator &eval,
     break;
 
   case ResolveWitnessResult::Missing: {
-    // The type witness is still missing. Resolve all of the type witnesses.
     auto &ctx = requirement->getASTContext();
+
+    if (ctx.LangOpts.EnableExperimentalAssociatedTypeInference) {
+      // Let's see if there is a better conformance we can perform associated
+      // type inference on.
+      auto *better = getBetterConformanceForResolvingTypeWitnesses(
+          conformance, requirement);
+
+      if (better != conformance &&
+          !ctx.evaluator.hasActiveRequest(ResolveTypeWitnessesRequest{better})) {
+        // Let's try to resolve type witnesses in the better conformance.
+        evaluateOrDefault(ctx.evaluator,
+                          ResolveTypeWitnessesRequest{better},
+                          evaluator::SideEffect());
+
+        // Check whether the above populated the type witness of our conformance.
+        auto known = conformance->TypeWitnesses.find(requirement);
+        if (known != conformance->TypeWitnesses.end())
+          return known->second;
+      }
+    }
+
+    // The type witness is still missing. Resolve all of the type witnesses
+    // in this conformance.
     evaluateOrDefault(ctx.evaluator,
                       ResolveTypeWitnessesRequest{conformance},
                       evaluator::SideEffect());

test/Generics/specialized_conformance_type_witness.swift

@@ -1,4 +1,5 @@
-// RUN: %target-typecheck-verify-swift
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -disable-experimental-associated-type-inference
 
 struct Row {}
 

test/decl/protocol/req/associated_type_inference_more_specific.swift

@@ -0,0 +1,32 @@
+// RUN: not %target-swift-frontend -typecheck %s -disable-experimental-associated-type-inference
+// RUN: %target-typecheck-verify-swift -enable-experimental-associated-type-inference
+
+protocol P {
+  associatedtype A
+}
+
+protocol Q: P {
+  associatedtype A
+  func f() -> A
+}
+
+protocol R: Q {}
+
+// We don't have enough information to infer 'A' just from the 'S1: P' and
+// 'S2: P' conformances. Make sure that if we force one of those conformances
+// first, we jump up to the conformance to Q, which has a requirement 'f()'
+// that gives us a way to infer 'A'.
+
+func forceP<T: P>(_: T) -> T.A {}
+
+let x: Int = forceP(S1())
+
+struct S1: Q {
+  func f() -> Int {}
+}
+
+let y: String = forceP(S2())
+
+struct S2: R {
+  func f() -> String {}
+}