[Sema] Improve diagnostic for closure return contextual mismatch

Author: LucianoPAlmeida

lib/Sema/CSApply.cpp

@@ -9124,8 +9124,9 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
       // Bodies of single-expression closures use a special locator
       // for contextual type conversion to make sure that result is
       // convertible to `Void` when `return` is not used explicitly.
-      if (contextualTypePurpose == CTP_ClosureResult) {
-        auto *closure = cast<ClosureExpr>(target.getDeclContext());
+      auto *closure = dyn_cast<ClosureExpr>(target.getDeclContext());
+      if (closure && closure->hasSingleExpressionBody() &&
+          contextualTypePurpose == CTP_ClosureResult) {
         auto *returnStmt =
             castToStmt<ReturnStmt>(closure->getBody()->getLastElement());
 

lib/Sema/CSClosure.cpp

@@ -891,8 +891,9 @@ class SyntacticElementConstraintGenerator
         return;
       }
 
+      auto contextualResultInfo = getContextualResultInfo();
       cs.addConstraint(ConstraintKind::Conversion, cs.getType(expr),
-                       getContextualResultType(),
+                       contextualResultInfo.getType(),
                        cs.getConstraintLocator(
                            context.getAbstractClosureExpr(),
                            LocatorPathElt::ClosureBody(
@@ -913,8 +914,10 @@ class SyntacticElementConstraintGenerator
       resultExpr = getVoidExpr(cs.getASTContext(), returnStmt->getEndLoc());
     }
 
+    auto contextualResultInfo = getContextualResultInfo();
     SolutionApplicationTarget target(resultExpr, context.getAsDeclContext(),
-                                     CTP_ReturnStmt, getContextualResultType(),
+                                     contextualResultInfo.purpose,
+                                     contextualResultInfo.getType(),
                                      /*isDiscarded=*/false);
 
     if (cs.generateConstraints(target, FreeTypeVariableBinding::Disallow)) {
@@ -923,8 +926,8 @@ class SyntacticElementConstraintGenerator
     }
 
     cs.setContextualType(target.getAsExpr(),
-                         TypeLoc::withoutLoc(getContextualResultType()),
-                         CTP_ReturnStmt);
+                         TypeLoc::withoutLoc(contextualResultInfo.getType()),
+                         contextualResultInfo.purpose);
     cs.setSolutionApplicationTarget(returnStmt, target);
   }
 
@@ -939,15 +942,15 @@ class SyntacticElementConstraintGenerator
     return context.hasSingleExpressionBody();
   }
 
-  Type getContextualResultType() const {
+  ContextualTypeInfo getContextualResultInfo() const {
     if (auto transform = cs.getAppliedResultBuilderTransform(context))
-      return transform->bodyResultType;
+      return {transform->bodyResultType, CTP_ReturnStmt};
 
     if (auto *closure =
             getAsExpr<ClosureExpr>(context.getAbstractClosureExpr()))
-      return cs.getClosureType(closure)->getResult();
+      return {cs.getClosureType(closure)->getResult(), CTP_ClosureResult};
 
-    return context.getBodyResultType();
+    return {context.getBodyResultType(), CTP_ReturnStmt};
   }
 
 #define UNSUPPORTED_STMT(STMT) void visit##STMT##Stmt(STMT##Stmt *) { \

test/Constraints/callAsFunction.swift

@@ -83,7 +83,7 @@ struct Test {
   var body8: MyLayout {
     MyLayout {
       let x = ""
-      return x // expected-error {{cannot convert return expression of type 'String' to return type 'Int'}}
+      return x // expected-error {{cannot convert value of type 'String' to closure result type 'Int'}}
     } content: {
       EmptyView()
     }

test/decl/circularity.swift

@@ -41,7 +41,7 @@ class Sub: Base {
     var foo = { () -> Int in
         let x = 42
         // FIXME: Bogus diagnostic
-        return foo(1) // expected-error {{cannot convert return expression of type '()' to return type 'Int'}}
+        return foo(1) // expected-error {{cannot convert value of type '()' to closure result type 'Int'}}
     }()
 }
 

test/expr/closure/closures.swift

@@ -690,7 +690,7 @@ func testSR13239_ArgsFn() -> Int {
 func testSR13239MultiExpr() -> Int {
   callit {
     print("hello") 
-    return print("hello") // expected-error {{cannot convert return expression of type '()' to return type 'Int'}}
+    return print("hello") // expected-error {{cannot convert value of type '()' to closure result type 'Int'}}
   }
 }
 
@@ -734,21 +734,20 @@ public class TestImplicitCaptureOfExplicitCaptureOfSelfInEscapingClosure {
 }
 
 // https://github.com/apple/swift/issues/59716
-// FIXME: Diagnostic should be tailored for closure result
 ["foo"].map { s in
-    if s == "1" { return } // expected-error{{cannot convert return expression of type '()' to return type 'Bool'}}
+    if s == "1" { return } // expected-error{{cannot convert value of type '()' to closure result type 'Bool'}}
     return s.isEmpty
 }.filter { $0 }
 
 ["foo"].map { s in
-    if s == "1" { return } // expected-error{{cannot convert return expression of type '()' to return type 'Bool'}}
+    if s == "1" { return } // expected-error{{cannot convert value of type '()' to closure result type 'Bool'}}
     if s == "2" { return }
     if s == "3" { return }
     return s.isEmpty
 }.filter { $0 }
 
 ["foo"].map { s in
-    if s == "1" { return () } // expected-error{{cannot convert return expression of type '()' to return type 'Bool'}}
+    if s == "1" { return () } // expected-error{{cannot convert value of type '()' to closure result type 'Bool'}}
     return s.isEmpty
 }.filter { $0 }