[Diagnostics] Port explicit closure result contextual mismatch

lib/Sema/CSDiag.cpp

@@ -3243,35 +3243,6 @@ bool FailureDiagnosis::diagnoseClosureExpr(
       return true;
   }
 
-  // If the body of the closure looked ok, then look for a contextual type
-  // error.  This is necessary because FailureDiagnosis::diagnoseExprFailure
-  // doesn't do this for closures.
-  if (contextualType) {
-    auto fnType = contextualType->getAs<AnyFunctionType>();
-    if (!fnType || fnType->isEqual(CS.getType(CE)))
-      return false;
-
-    auto contextualResultType = fnType->getResult();
-    // If the result type was unknown, it doesn't really make
-    // sense to diagnose from expected to unknown here.
-    if (isInvalidClosureResultType(contextualResultType))
-      return false;
-
-    // If the closure had an explicitly written return type incompatible with
-    // the contextual type, diagnose that.
-    if (CE->hasExplicitResultType() &&
-        CE->getExplicitResultTypeLoc().getTypeRepr()) {
-      auto explicitResultTy = CE->getExplicitResultTypeLoc().getType();
-      if (fnType && !explicitResultTy->isEqual(contextualResultType)) {
-        auto repr = CE->getExplicitResultTypeLoc().getTypeRepr();
-        diagnose(repr->getStartLoc(), diag::incorrect_explicit_closure_result,
-                 explicitResultTy, fnType->getResult())
-          .fixItReplace(repr->getSourceRange(),fnType->getResult().getString());
-        return true;
-      }
-    }
-  }
-
   // Otherwise, we can't produce a specific diagnostic.
   return false;
 }

lib/Sema/CSDiagnostics.cpp

@@ -1986,6 +1986,17 @@ bool ContextualFailure::diagnoseAsError() {
   Diag<Type, Type> diagnostic;
   switch (path.back().getKind()) {
   case ConstraintLocator::ClosureResult: {
+    auto *closure = cast<ClosureExpr>(getRawAnchor());
+    if (closure->hasExplicitResultType() &&
+        closure->getExplicitResultTypeLoc().getTypeRepr()) {
+      auto resultRepr = closure->getExplicitResultTypeLoc().getTypeRepr();
+      emitDiagnostic(resultRepr->getStartLoc(),
+                     diag::incorrect_explicit_closure_result, getFromType(),
+                     getToType())
+          .fixItReplace(resultRepr->getSourceRange(), getToType().getString());
+      return true;
+    }
+
     diagnostic = diag::cannot_convert_closure_result;
     break;
   }

lib/Sema/CSSimplify.cpp

@@ -3265,8 +3265,25 @@ bool ConstraintSystem::repairFailures(
     break;
   }
 
+  case ConstraintLocator::FunctionResult: {
+    auto *loc = getConstraintLocator(anchor, {path.begin(), path.end() - 1});
+    // If this is a mismatch between contextual type and (trailing)
+    // closure with explicitly specified result type let's record it
+    // as contextual type mismatch.
+    if (loc->isLastElement<LocatorPathElt::ContextualType>() ||
+        loc->isLastElement<LocatorPathElt::ApplyArgToParam>()) {
+      auto *argExpr = simplifyLocatorToAnchor(loc);
+      if (argExpr && isa<ClosureExpr>(argExpr)) {
+        conversionsOrFixes.push_back(ContextualMismatch::create(
+            *this, lhs, rhs,
+            getConstraintLocator(argExpr, ConstraintLocator::ClosureResult)));
+        break;
+      }
+    }
+    LLVM_FALLTHROUGH;
+  }
+
   case ConstraintLocator::Member:
-  case ConstraintLocator::FunctionResult:
   case ConstraintLocator::DynamicLookupResult: {
     // Most likely this is an attempt to use get-only subscript as mutating,
     // or assign a value of a result of function/member ref e.g. `foo() = 42`

test/Constraints/closures.swift

@@ -765,7 +765,7 @@ overloaded { print("hi"); print("bye") } // multiple expression closure without
 func not_overloaded(_ handler: () -> Int) {}
 
 not_overloaded { } // empty body
-// expected-error@-1 {{cannot convert value of type '() -> ()' to expected argument type '() -> Int'}}
+// expected-error@-1 {{cannot convert value of type '()' to closure result type 'Int'}}
 
 not_overloaded { print("hi") } // single-expression closure
 // expected-error@-1 {{cannot convert value of type '()' to closure result type 'Int'}}
@@ -786,7 +786,7 @@ func test() -> Int? {
 }
 
 var fn: () -> [Int] = {}
-// expected-error@-1 {{cannot convert value of type '() -> ()' to specified type '() -> [Int]'}}
+// expected-error@-1 {{cannot convert value of type '()' to closure result type '[Int]'}}
 
 fn = {}
 // expected-error@-1 {{cannot assign value of type '() -> ()' to type '() -> [Int]'}}

test/Sema/immutability.swift

@@ -656,8 +656,8 @@ func sr4214() {
   let closure = { val in val.x = 7 } as (inout MutableSubscripts) -> () // Ok
   var v = MutableSubscripts()
   closure(&v)
-  // FIXME: This diagnostic isn't really all that much better
-  // expected-error@+1 {{cannot convert value of type '(inout MutableSubscripts) -> ()' to expected argument type '(_) -> _'}}
+  // expected-error@+2 {{declared closure result '()' is incompatible with contextual type 'MutableSubscripts'}}
+  // expected-error@+1 {{cannot convert value of type '(inout MutableSubscripts) -> ()' to expected argument type '(MutableSubscripts) -> MutableSubscripts'}}
   sequence(v) { (state : inout MutableSubscripts) -> () in
     _ = MutableSubscripts.initialize(from: &state)
     return ()

test/Sema/substring_to_string_conversion_swift4.swift

@@ -45,7 +45,7 @@ do {
 
 // CTP_ClosureResult
 do {
-  [ss].map { (x: Substring) -> String in x } // expected-error {{cannot convert value of type 'Substring' to closure result type 'String'}} {{42-42=String(}} {{43-43=)}}
+  [ss].map { (x: Substring) -> String in x } // expected-error {{declared closure result 'Substring' is incompatible with contextual type 'String'}}
 }
 
 // CTP_ArrayElement