[Sema] [Diagnostics] Consider function result locator when get structural contextual type for function arg/param

Author: LucianoPAlmeida

lib/Sema/CSFix.cpp

@@ -403,9 +403,21 @@ getStructuralTypeContext(const Solution &solution, ConstraintLocator *locator) {
     return std::make_tuple(contextualTypeElt->getPurpose(), exprType,
                            contextualType);
   } else if (auto argApplyInfo = solution.getFunctionArgApplyInfo(locator)) {
-    return std::make_tuple(CTP_CallArgument,
-                           argApplyInfo->getArgType(),
-                           argApplyInfo->getParamType());
+    Type fromType = argApplyInfo->getArgType();
+    Type toType = argApplyInfo->getParamType();
+    // In case locator points to the function result we want the
+    // argument and param function types result.
+    if (locator->isLastElement<LocatorPathElt::FunctionResult>()) {
+      auto fromFnType = fromType->getAs<FunctionType>();
+      auto toFnType = toType->getAs<FunctionType>();
+      if (fromFnType && toFnType) {
+        auto &cs = solution.getConstraintSystem();
+        return std::make_tuple(
+            cs.getContextualTypePurpose(locator->getAnchor()),
+            fromFnType->getResult(), toFnType->getResult());
+      }
+    }
+    return std::make_tuple(CTP_CallArgument, fromType, toType);
   } else if (auto *coerceExpr = getAsExpr<CoerceExpr>(locator->getAnchor())) {
     return std::make_tuple(CTP_CoerceOperand,
                            solution.getType(coerceExpr->getSubExpr()),

test/expr/closure/closures.swift

@@ -589,3 +589,20 @@ class SR14120 {
     }
   }
 }
+
+// SR-14678
+func call<T>(_ : Int, _ f: () -> (T, Int)) -> (T, Int) {
+  f()
+}
+
+func testSR14678() -> (Int, Int) {
+  call(1) { // expected-error {{cannot convert return expression of type '((), Int)' to return type '(Int, Int)'}}
+     (print("hello"), 0)
+  }
+}
+
+func testSR14678_Optional() -> (Int, Int)? {
+  call(1) { // expected-error {{cannot convert return expression of type '((), Int)' to return type '(Int, Int)'}}
+     (print("hello"), 0)
+  }
+}