[ConstraintSystem] Retrieve contextual type from a solution for ambiguities

In ambiguity scenarios solutions are not applied back to the constraint
system, so it might not always have contextual type information when it
was recorded e.g. for a multi-statement closure.

Resolves: rdar://97396399

(cherry picked from commit b07c887)
(cherry picked from commit 399d349)
(cherry picked from commit c300edb)

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -267,6 +267,10 @@ ERROR(no_candidates_match_result_type,none,
       "no '%0' candidates produce the expected contextual result type %1",
       (StringRef, Type))
 
+ERROR(no_overloads_have_result_type_conformance,none,
+      "no '%0' overloads produce result type that conforms to %1",
+      (StringRef, Type))
+
 ERROR(no_candidates_match_argument_type,none,
       "no '%0' candidates produce the expected type %1 for parameter #%2",
       (StringRef, Type, unsigned))
@@ -559,6 +563,10 @@ NOTE(cannot_convert_candidate_result_to_contextual_type,none,
      "%0 produces %1, not the expected contextual result type %2",
      (DeclName, Type, Type))
 
+NOTE(overload_result_type_does_not_conform,none,
+     "result type %1 of %0 does not conform to %2",
+     (DeclName, Type, Type))
+
 // for ... in expression
 ERROR(cannot_convert_sequence_element_value,none,
       "cannot convert sequence element type %0 to expected type %1",

include/swift/Sema/ConstraintSystem.h

@@ -1422,6 +1422,15 @@ class Solution {
   /// types where possible.
   Type resolveInterfaceType(Type type) const;
 
+  Type getContextualType(ASTNode anchor) const {
+    for (const auto &entry : contextualTypes) {
+      if (entry.first == anchor) {
+        return simplifyType(entry.second.getType());
+      }
+    }
+    return Type();
+  }
+
   /// For a given locator describing a function argument conversion, or a
   /// constraint within an argument conversion, returns information about the
   /// application of the argument to its parameter. If the locator is not

lib/Sema/ConstraintSystem.cpp

@@ -4219,27 +4219,31 @@ static bool diagnoseAmbiguityWithContextualType(
 
   auto anchor = locator->getAnchor();
   auto name = result->choices.front().getName();
-  DE.diagnose(getLoc(anchor), diag::no_candidates_match_result_type,
-              name.getBaseName().userFacingName(),
-              cs.getContextualType(anchor, /*forConstraint=*/false));
+  auto contextualTy = solution.getContextualType(anchor);
+
+  assert(contextualTy);
+
+  DE.diagnose(getLoc(anchor),
+              contextualTy->is<ProtocolType>()
+                  ? diag::no_overloads_have_result_type_conformance
+                  : diag::no_candidates_match_result_type,
+              name.getBaseName().userFacingName(), contextualTy);
 
   for (const auto &solution : solutions) {
     auto overload = solution.getOverloadChoice(calleeLocator);
     if (auto *decl = overload.choice.getDeclOrNull()) {
-      auto loc = decl->getLoc();
-      if (loc.isInvalid())
-        continue;
-
       auto type = solution.simplifyType(overload.boundType);
 
       if (isExpr<ApplyExpr>(anchor) || isExpr<SubscriptExpr>(anchor)) {
         auto fnType = type->castTo<FunctionType>();
         DE.diagnose(
-            loc, diag::cannot_convert_candidate_result_to_contextual_type,
-            decl->getName(), fnType->getResult(),
-            cs.getContextualType(anchor, /*forConstraint=*/false));
+            decl,
+            contextualTy->is<ProtocolType>()
+                ? diag::overload_result_type_does_not_conform
+                : diag::cannot_convert_candidate_result_to_contextual_type,
+            decl->getName(), fnType->getResult(), contextualTy);
       } else {
-        DE.diagnose(loc, diag::found_candidate_type, type);
+        DE.diagnose(decl, diag::found_candidate_type, type);
       }
     }
   }

test/Constraints/overload.swift

@@ -257,3 +257,30 @@ func rdar79672230() {
   var t: MyType = MyType()
   test(&t) // expected-error {{no exact matches in call to local function 'test'}}
 }
+
+// rdar://97396399 - crash in swift::DiagnosticEngine::formatDiagnosticText
+func rdar97396399() {
+  // Has to be overloaded to make sure that contextual type is not recorded during constraint generation
+  func test(_: () -> Void) {}
+  func test(_: (Int) -> Void) {}
+
+  // Multiple different overloads, none of which conform to Sequence
+  func fn(_: Int) -> Int {}
+  // expected-note@-1 {{found candidate with type '(Int) -> Int'}}
+  // expected-note@-2 {{result type 'Int' of 'fn' does not conform to 'Sequence'}}
+  func fn(_: Int) -> Double {}
+  // expected-note@-1 {{found candidate with type '(Int) -> Double'}}
+  // expected-note@-2 {{result type 'Double' of 'fn' does not conform to 'Sequence'}}
+
+  test {
+    for x in fn { // expected-error {{no 'fn' overloads produce result type that conforms to 'Sequence'}}
+      print(x)
+    }
+  }
+
+  test {
+    for x in fn(42) { // expected-error {{no 'fn' overloads produce result type that conforms to 'Sequence'}}
+      print(x)
+    }
+  }
+}