[ConstraintSystem] Don't use special locator for value-to-value conversions related to contextual type

Previously locator for value-to-value conversion would just drop
all the contextual information if the conversion occurred while
converting expression to a contextual type. This is incorrect for
i.e. `return` statements and other targets because because they
are solved separately and using the same locator would result in
a clash when solutions are merged.

Author: xedin

lib/Sema/CSApply.cpp

@@ -6792,12 +6792,25 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
     case ConversionRestrictionKind::DoubleToCGFloat: {
       auto conversionKind = knownRestriction->second;
 
-      // If conversion wraps the whole body of a single-expression closure,
-      // let's use the passed-in expression since the closure itself doesn't
-      // get updated until coercion is done.
-      auto *argExpr = locator.endsWith<LocatorPathElt::ClosureBody>()
-                          ? expr
-                          : locator.trySimplifyToExpr();
+      auto shouldUseCoercedExpr = [&]() {
+        // If conversion wraps the whole body of a single-expression closure,
+        // let's use the passed-in expression since the closure itself doesn't
+        // get updated until coercion is done.
+        if (locator.endsWith<LocatorPathElt::ClosureBody>())
+          return true;
+
+        // Contextual type locator always uses the original version of
+        // expression (before any coercions have been applied) because
+        // otherwise it wouldn't be possible to find the overload choice.
+        if (locator.endsWith<LocatorPathElt::ContextualType>())
+          return true;
+
+        // In all other cases use the expression associated with locator.
+        return false;
+      };
+
+      auto *argExpr =
+          shouldUseCoercedExpr() ? expr : locator.trySimplifyToExpr();
       assert(argExpr);
 
       // Source requires implicit conversion to match destination
@@ -8575,7 +8588,7 @@ static Optional<SolutionApplicationTarget> applySolutionToInitialization(
   // Convert the initializer to the type of the pattern.
   auto &cs = solution.getConstraintSystem();
   auto locator = cs.getConstraintLocator(
-      initializer, LocatorPathElt::ContextualType(CTP_Initialization));
+      target.getAsExpr(), LocatorPathElt::ContextualType(CTP_Initialization));
   initializer = solution.coerceToType(initializer, initType, locator);
   if (!initializer)
     return None;
@@ -9087,10 +9100,12 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     // do so now.
     if (shouldCoerceToContextualType()) {
       ConstraintLocator *locator = nullptr;
+
+      auto contextualTypePurpose = target.getExprContextualTypePurpose();
       // 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 (target.getExprContextualTypePurpose() == CTP_ClosureResult) {
+      if (contextualTypePurpose == CTP_ClosureResult) {
         auto *closure = cast<ClosureExpr>(target.getDeclContext());
         auto *returnStmt =
             castToStmt<ReturnStmt>(closure->getBody()->getLastElement());
@@ -9100,8 +9115,7 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
                          /*hasReturn=*/!returnStmt->isImplicit()));
       } else {
         locator = cs.getConstraintLocator(
-            resultExpr, LocatorPathElt::ContextualType(
-                            target.getExprContextualTypePurpose()));
+            expr, LocatorPathElt::ContextualType(contextualTypePurpose));
       }
 
       assert(locator);
@@ -9113,7 +9127,7 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
       // We referenced an lvalue. Load it.
       resultExpr = Rewriter.coerceToType(
           resultExpr, cs.getType(resultExpr)->getRValueType(),
-          cs.getConstraintLocator(resultExpr,
+          cs.getConstraintLocator(expr,
                                   LocatorPathElt::ContextualType(
                                       target.getExprContextualTypePurpose())));
     }

lib/Sema/ConstraintSystem.cpp

@@ -444,15 +444,6 @@ ConstraintLocator *ConstraintSystem::getImplicitValueConversionLocator(
       break;
     }
 
-    // If the conversion is associated with a contextual type e.g.
-    // `_: Double = CGFloat(1)` then drop `ContextualType` so that
-    // it's easy to find when the underlying expression has been
-    // rewritten.
-    if (!path.empty() && path.back().is<LocatorPathElt::ContextualType>()) {
-      anchor = ASTNode();
-      path.clear();
-    }
-
     // If conversion is for a tuple element, let's drop `TupleType`
     // components from the path since they carry information for
     // diagnostics that `ExprRewriter` won't be able to re-construct

test/Constraints/implicit_double_cgfloat_conversion.swift

@@ -285,3 +285,33 @@ func test_conversion_inside_tuple_elements() -> (a: CGFloat, b: (c: Int, d: CGFl
   let x: Double = 0.0
   return (a: x, b: (c: 42, d: x)) // Ok
 }
+
+do {
+  struct Data {
+    var prop: CGFloat
+  }
+
+  func single(get: () -> Double) {}
+  func multiple(get1: () -> Double,
+                get2: () -> CGFloat = { Double(1) },
+                get3: () -> Double) {}
+
+  func test(data: Data) {
+    single { data.prop } // Ok
+    single { return data.prop } // Ok
+
+    single {
+      _ = 42
+      if true {
+        return data.prop // Ok
+      }
+      return data.prop // Ok
+    }
+
+    multiple {
+      data.prop // Ok
+    } get3: {
+      return data.prop // Ok
+    }
+  }
+}