[CSBindings] Fix inference of partial key path from conversion constraints

It's not permitted to use `PartialKeyPath` type to resolve a key path
expression, because that type is intended to be a type-erased version of
a fully resolved `KeyPath` type.

In situations where contextual type is a partial key path (e.g. parameter
type is a partial key path), let's replace it with a `KeyPath` type that
is a subtype of `PartialKeyPath` and allow value inference to happen.

Resolves: SR-5667
Resolves: rdar://32365183

Author: xedin

include/swift/Sema/ConstraintLocator.h

@@ -885,6 +885,22 @@ class LocatorPathElt::ContextualType final : public StoredIntegerElement<1> {
   }
 };
 
+class LocatorPathElt::KeyPathType final
+    : public StoredPointerElement<TypeBase> {
+public:
+  KeyPathType(Type valueType)
+      : StoredPointerElement(PathElementKind::KeyPathType,
+                             valueType.getPointer()) {
+    assert(valueType);
+  }
+
+  Type getValueType() const { return getStoredPointer(); }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == PathElementKind::KeyPathType;
+  }
+};
+
 namespace details {
   template <typename CustomPathElement>
   class PathElement {

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -109,7 +109,7 @@ CUSTOM_LOCATOR_PATH_ELT(KeyPathDynamicMember)
 SIMPLE_LOCATOR_PATH_ELT(KeyPathRoot)
 
 /// The type of the key path expression.
-SIMPLE_LOCATOR_PATH_ELT(KeyPathType)
+CUSTOM_LOCATOR_PATH_ELT(KeyPathType)
 
 /// The value of a key path.
 SIMPLE_LOCATOR_PATH_ELT(KeyPathValue)

include/swift/Sema/ConstraintSystem.h

@@ -356,6 +356,10 @@ class TypeVariableType::Implementation {
   /// Determine whether this type variable represents a closure result type.
   bool isClosureResultType() const;
 
+  /// Determine whether this type variable represents
+  /// a type of a key path expression.
+  bool isKeyPathType() const;
+
   /// Retrieve the representative of the equivalence class to which this
   /// type variable belongs.
   ///

lib/Sema/CSBindings.cpp

@@ -1110,14 +1110,36 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
   if (type->hasError())
     return None;
 
-  if (auto *locator = TypeVar->getImpl().getLocator()) {
-    if (locator->isKeyPathType()) {
-      auto *BGT =
-          type->lookThroughAllOptionalTypes()->getAs<BoundGenericType>();
-      if (!BGT || !isKnownKeyPathDecl(CS.getASTContext(), BGT->getDecl()))
-        return None;
+  if (TypeVar->getImpl().isKeyPathType()) {
+    auto *BGT = type->lookThroughAllOptionalTypes()->getAs<BoundGenericType>();
+    if (!BGT || !isKnownKeyPathDecl(CS.getASTContext(), BGT->getDecl()))
+      return None;
+
+    // `PartialKeyPath<T>` represents a type-erased version of `KeyPath<T, V>`.
+    //
+    // In situations where partial key path cannot be used directly i.e.
+    // passing an argument to a parameter represented by a partial key path,
+    // let's attempt a `KeyPath` binding which would then be converted to a
+    // partial key path since there is a subtype relationship between them.
+    auto &ctx = CS.getASTContext();
+    if (BGT->getDecl() == ctx.getPartialKeyPathDecl() &&
+        kind == AllowedBindingKind::Subtypes) {
+      auto *keyPathLoc = TypeVar->getImpl().getLocator();
+
+      auto rootTy = BGT->getGenericArgs()[0];
+      // Since partial key path is an erased version of `KeyPath`, the value
+      // type would never be used, which means that binding can use
+      // type variable generated for a result of key path expression.
+      auto valueTy =
+          keyPathLoc->castLastElementTo<LocatorPathElt::KeyPathType>()
+              .getValueType();
+
+      type = BoundGenericType::get(ctx.getKeyPathDecl(), Type(),
+                                   {rootTy, valueTy});
     }
+  }
 
+  if (auto *locator = TypeVar->getImpl().getLocator()) {
     // Don't allow a protocol type to get propagated from the base to the result
     // type of a chain, Result should always be a concrete type which conforms
     // to the protocol inferred for the base.

lib/Sema/CSGen.cpp

@@ -3257,10 +3257,11 @@ namespace {
       // The result is a KeyPath from the root to the end component.
       // The type of key path depends on the overloads chosen for the key
       // path components.
-      auto typeLoc =
-          CS.getConstraintLocator(locator, ConstraintLocator::KeyPathType);
+      auto typeLoc = CS.getConstraintLocator(
+          locator, LocatorPathElt::KeyPathType(rvalueBase));
+
       Type kpTy = CS.createTypeVariable(typeLoc, TVO_CanBindToNoEscape |
-                                        TVO_CanBindToHole);
+                                                     TVO_CanBindToHole);
       CS.addKeyPathConstraint(kpTy, root, rvalueBase, componentTypeVars,
                               locator);
       return kpTy;

lib/Sema/ConstraintLocator.cpp

@@ -121,7 +121,7 @@ bool ConstraintLocator::isKeyPathType() const {
   // The format of locator should be `<keypath expr> -> key path type`
   if (!anchor || !isExpr<KeyPathExpr>(anchor) || path.size() != 1)
     return false;
-  return path.back().getKind() == ConstraintLocator::KeyPathType;
+  return path.back().is<LocatorPathElt::KeyPathType>();
 }
 
 bool ConstraintLocator::isKeyPathRoot() const {

lib/Sema/TypeCheckConstraints.cpp

@@ -111,6 +111,10 @@ bool TypeVariableType::Implementation::isClosureResultType() const {
          locator->isLastElement<LocatorPathElt::ClosureResult>();
 }
 
+bool TypeVariableType::Implementation::isKeyPathType() const {
+  return locator && locator->isKeyPathType();
+}
+
 void *operator new(size_t bytes, ConstraintSystem& cs,
                    size_t alignment) {
   return cs.getAllocator().Allocate(bytes, alignment);

test/expr/unary/keypath/keypath.swift

@@ -180,14 +180,10 @@ func testKeyPath(sub: Sub, optSub: OptSub,
   var m = [\A.property, \A.[sub], \A.optProperty!]
   expect(&m, toHaveType: Exactly<[PartialKeyPath<A>]>.self)
 
-  // \.optProperty returns an optional of Prop and `\.[sub]` returns `A`
-  // expected-error@+2 {{key path value type 'Prop?' cannot be converted to contextual type 'Prop'}}
-  // expected-error@+1 {{key path value type 'A' cannot be converted to contextual type 'Prop'}}
+  // \.optProperty returns an optional of Prop and `\.[sub]` returns `A`, all this unifies into `[PartialKeyPath<A>]`
   var n = [\A.property, \.optProperty, \.[sub], \.optProperty!]
   expect(&n, toHaveType: Exactly<[PartialKeyPath<A>]>.self)
 
-  // FIXME: shouldn't be ambiguous
-  // expected-error@+1{{ambiguous}}
   let _: [PartialKeyPath<A>] = [\.property, \.optProperty, \.[sub], \.optProperty!]
 
   var o = [\A.property, \C<A>.value]
@@ -1111,3 +1107,26 @@ func test_kp_as_function_mismatch() {
   _ = a.filter(\String.filterOut) // expected-error{{key path value type '(String) throws -> Bool' cannot be converted to contextual type 'Bool'}}
 
 }
+
+func test_partial_keypath_inference() {
+  // rdar://problem/34144827
+
+  struct S { var i: Int = 0 }
+  enum E { case A(pkp: PartialKeyPath<S>) }
+
+  _ = E.A(pkp: \.i) // Ok
+
+  // rdar://problem/36472188
+
+  class ThePath {
+    var isWinding:Bool?
+  }
+
+  func walk<T>(aPath: T, forKey: PartialKeyPath<T>) {}
+  func walkThePath(aPath: ThePath, forKey: PartialKeyPath<ThePath>) {}
+
+  func test(path: ThePath) {
+    walkThePath(aPath: path, forKey: \.isWinding) // Ok
+    walk(aPath: path, forKey: \.isWinding) // Ok
+  }
+}