[CSSimplify] Prevent overly eager solving of key path constraint

If key path literal is passed as an argument to a function/subscript
application its type cannot be resolved until the overload for the
call is selected, otherwise it could prevent a valid keypath-to-function
conversion in cases were argument ends up being a function type.

Author: xedin

lib/Sema/CSSimplify.cpp

@@ -12364,21 +12364,56 @@ ConstraintSystem::simplifyKeyPathConstraint(
       kpDecl = getASTContext().getWritableKeyPathDecl();
   }
 
+  auto formUnsolved = [&]() {
+    addUnsolvedConstraint(Constraint::create(
+        *this, ConstraintKind::KeyPath, keyPathTy, rootTy, valueTy,
+        locator.getBaseLocator(), componentTypeVars));
+  };
+
   auto loc = locator.getBaseLocator();
   if (definitelyFunctionType) {
     increaseScore(SK_FunctionConversion, locator);
     return SolutionKind::Solved;
   } else if (!anyComponentsUnresolved ||
              (definitelyKeyPathType && capability == ReadOnly)) {
-    auto resolvedKPTy =
-      BoundGenericType::get(kpDecl, nullptr, {rootTy, valueTy});
-    return matchTypes(resolvedKPTy, keyPathTy, ConstraintKind::Bind, subflags,
-                      loc);
+    // If key path is connected to a function application it cannot be
+    // bound until the choice is selected because it's undetermined
+    // until then whether key path is implicitly converted to a function
+    // type or not.
+    //
+    // \code
+    // struct Data {
+    //   var value: Int = 42
+    // }
+    //
+    // func test<S: Sequence>(_: S, _: (S.Element) -> Int) {}
+    // func test<C: Collection>(_: C, _: (C.Element) -> Int) {}
+    //
+    // func test(arr: [Data]) {
+    //  test(arr, \Data.value)
+    // }
+    // \endcode
+    //
+    // In this example if we did allow to bind the key path before
+    // an overload of `test` is selected, we'd end up with no solutions
+    // because the type of the key path expression is actually: '(Data) -> Int'
+    // and not 'WritableKeyPath<Data, Int>`.
+    auto *typeVar = keyPathTy->getAs<TypeVariableType>();
+    if (typeVar && findConstraintThroughOptionals(
+                       typeVar, OptionalWrappingDirection::Unwrap,
+                       [&](Constraint *match, TypeVariableType *) {
+                         return match->getKind() ==
+                                ConstraintKind::ApplicableFunction;
+                       })) {
+      formUnsolved();
+    } else {
+      auto resolvedKPTy =
+          BoundGenericType::get(kpDecl, nullptr, {rootTy, valueTy});
+      return matchTypes(resolvedKPTy, keyPathTy, ConstraintKind::Bind, subflags,
+                        loc);
+    }
   } else {
-    addUnsolvedConstraint(Constraint::create(*this, ConstraintKind::KeyPath,
-                                             keyPathTy, rootTy, valueTy,
-                                             locator.getBaseLocator(),
-                                             componentTypeVars));
+    formUnsolved();
   }
   return SolutionKind::Solved;
 }

test/expr/unary/keypath/keypath.swift

@@ -1202,3 +1202,17 @@ func test_keypath_inference_from_existentials() {
   test(\.other, 42)  // Ok
   test(\.member, "") // Ok
 }
+
+// rdar://116376651 - key path type is bound before context is fully resolved.
+func keypath_to_func_conversion_as_arg_to_overloaded_func() {
+  struct Data {
+    var value: Int = 42
+  }
+
+  func test<S: Sequence>(_: S, _: (S.Element) -> Int) {}
+  func test<C: Collection>(_: C, _: (C.Element) -> Int) {}
+
+  func test(arr: [Data]) {
+    test(arr, \Data.value) // Ok
+  }
+}

test/expr/unary/keypath/salvage-with-other-type-errors.swift

@@ -6,7 +6,7 @@
 struct P<T: K> { }
 
 struct S {
-    init<B>(_ a: P<B>) { // expected-note {{in call to initializer}}
+    init<B>(_ a: P<B>) { // expected-note {{where 'B' = 'String'}}
         fatalError()
     }
 }
@@ -27,9 +27,8 @@ struct A {
 }
 
 extension A: K {
-  static let j = S(\A.id + "id") // expected-error {{generic parameter 'B' could not be inferred}}
-  // expected-error@-1 {{binary operator '+' cannot be applied to operands of type 'KeyPath<A, String>' and 'String'}}
-  // expected-note@-2 {{overloads for '+' exist with these partially matching parameter lists: (String, String)}}
+  static let j = S(\A.id + "id")
+  // expected-error@-1 {{initializer 'init(_:)' requires that 'String' conform to 'K'}}
 }
 
 // https://github.com/apple/swift/issues/47610