[ConstraintSystem] Adjust contextual locators to support updated Double/CGFloat conversion

include/swift/Sema/ConstraintLocator.h

@@ -1295,6 +1295,16 @@ class ConstraintLocatorBuilder {
     return None;
   }
 
+  /// Check whether this locator has the given locator path element
+  /// at the end of its path.
+  template <class Kind>
+  bool endsWith() const {
+    if (auto lastElt = last()) {
+      return lastElt->is<Kind>();
+    }
+    return false;
+  }
+
   /// Produce a debugging dump of this locator.
   SWIFT_DEBUG_DUMPER(dump(SourceManager *SM));
   SWIFT_DEBUG_DUMPER(dump(ConstraintSystem *CS));

lib/Sema/CSApply.cpp

@@ -6792,7 +6792,25 @@ Expr *ExprRewriter::coerceToType(Expr *expr, Type toType,
     case ConversionRestrictionKind::DoubleToCGFloat: {
       auto conversionKind = knownRestriction->second;
 
-      auto *argExpr = 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
@@ -8570,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;
@@ -9081,17 +9099,35 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
     // If we're supposed to convert the expression to some particular type,
     // 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 (contextualTypePurpose == CTP_ClosureResult) {
+        auto *closure = cast<ClosureExpr>(target.getDeclContext());
+        auto *returnStmt =
+            castToStmt<ReturnStmt>(closure->getBody()->getLastElement());
+
+        locator = cs.getConstraintLocator(
+            closure, LocatorPathElt::ClosureBody(
+                         /*hasReturn=*/!returnStmt->isImplicit()));
+      } else {
+        locator = cs.getConstraintLocator(
+            expr, LocatorPathElt::ContextualType(contextualTypePurpose));
+      }
+
+      assert(locator);
+
       resultExpr = Rewriter.coerceToType(
-          resultExpr, solution.simplifyType(convertType),
-          cs.getConstraintLocator(resultExpr,
-                                  LocatorPathElt::ContextualType(
-                                      target.getExprContextualTypePurpose())));
+          resultExpr, solution.simplifyType(convertType), locator);
     } else if (cs.getType(resultExpr)->hasLValueType() &&
                !target.isDiscardedExpr()) {
       // 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/CSClosure.cpp

@@ -1662,7 +1662,7 @@ class SyntacticElementSolutionApplication
       assert(isSingleExpression);
       resultTarget = SolutionApplicationTarget(
           resultExpr, context.getAsDeclContext(),
-          mode == convertToResult ? CTP_ReturnStmt : CTP_Unused,
+          mode == convertToResult ? CTP_ClosureResult : CTP_Unused,
           mode == convertToResult ? resultType : Type(),
           /*isDiscarded=*/false);
     }

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
+    }
+  }
+}