[Diagnostics] Contextual mismatch should be attached to the closure if its body is empty/implicit

If there is a contextual mismatch associated with a closure body,
make sure that the diagnostic is attached to the closure even
if the body is empty or implicit.

Resolves: rdar://52204608

Author: xedin

lib/Sema/CSDiagnostics.cpp

@@ -2121,6 +2121,24 @@ Diag<StringRef> AssignmentFailure::findDeclDiagonstic(ASTContext &ctx,
   return diag::assignment_lhs_is_immutable_variable;
 }
 
+SourceLoc ContextualFailure::getLoc() const {
+  auto *locator = getLocator();
+
+  // `getSingleExpressionBody` can point to an implicit expression
+  // without source information in cases like `{ return }`.
+  if (locator->isLastElement<LocatorPathElt::ClosureBody>()) {
+    auto *closure = castToExpr<ClosureExpr>(locator->getAnchor());
+    if (closure->hasSingleExpressionBody()) {
+      auto *body = closure->getSingleExpressionBody();
+      if (auto loc = body->getLoc())
+        return loc;
+    }
+    return closure->getLoc();
+  }
+
+  return FailureDiagnostic::getLoc();
+}
+
 bool ContextualFailure::diagnoseAsError() {
   auto anchor = getAnchor();
   auto path = getLocator()->getPath();

lib/Sema/CSDiagnostics.h

@@ -603,6 +603,8 @@ class ContextualFailure : public FailureDiagnostic {
       : FailureDiagnostic(solution, locator), CTP(purpose), RawFromType(lhs),
         RawToType(rhs) {}
 
+  SourceLoc getLoc() const override;
+
   Type getFromType() const { return resolve(RawFromType); }
 
   Type getToType() const { return resolve(RawToType); }

test/Constraints/closures.swift

@@ -1077,3 +1077,10 @@ let _: (@convention(c) () -> Void)? = Bool.random() ? nil : {} // OK on type che
 let _: (@convention(block) () -> Void)? = Bool.random() ? {} : {} // OK
 let _: (@convention(thin) () -> Void)? = Bool.random() ? {} : {} // OK
 let _: (@convention(c) () -> Void)? = Bool.random() ? {} : {} // OK on type checking, diagnostics are deffered to SIL
+
+// Make sure that diagnostic is attached to the closure even when body is empty (implicitly returns `Void`)
+var emptyBodyMismatch: () -> Int {
+  return { // expected-error {{cannot convert value of type '()' to closure result type 'Int'}}
+    return
+  }
+}