[CSSimplify] Adjust {Any, Partial}KeyPath bindings right before resolving

A type variable that represents a key path literal cannot be bound
directly to `AnyKeyPath` or `PartialKeyPath`, such types could only
be used for conversions.
It used to be the task of the binding inference to infer `AnyKeyPath`
and `PartiaKeyPath` as a `KeyPath` using previously generated type
variables for a root and value associated with key path literal
(previously stored in the locator).

Recently we switched over to storing key path information in the
constraint system and introduced `resolveKeyPath` method to gain
more control over how key path type variable gets assigned.

Getting information from the constraint system creates a problem
for the inference because "undo" for some bindings would be run
after solver scope has been erased, so instead of modifying bindings
in `inferFromRelational`, let's do that in `resolveKeyPath` right
before the key path type variable gets bound which seems to be a
better place for that logic anyway.

Resolves: rdar://113760727

Author: xedin

lib/Sema/CSBindings.cpp

@@ -1253,35 +1253,6 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
     auto objectTy = type->lookThroughAllOptionalTypes();
     if (!isKnownKeyPathType(objectTy))
       return llvm::None;
-    
-    auto &ctx = CS.getASTContext();
-    auto *keyPathTypeLoc = TypeVar->getImpl().getLocator();
-    auto *keyPath = castToExpr<KeyPathExpr>(keyPathTypeLoc->getAnchor());
-    // `AnyKeyPath` and `PartialKeyPath` represent type-erased versions of
-    // `KeyPath<T, V>`.
-    //
-    // In situations where `AnyKeyPath` or `PartialKeyPath` cannot be used
-    // directly i.e. passing an argument to a parameter represented by a
-    // `AnyKeyPath` or `PartialKeyPath`, let's attempt a `KeyPath` binding which
-    // would then be converted to a `AnyKeyPath` or `PartialKeyPath` since there
-    // is a subtype relationship between them.
-    if (objectTy->isAnyKeyPath()) {
-      auto root = CS.getKeyPathRootType(keyPath);
-      auto value = CS.getKeyPathValueType(keyPath);
-      
-      type = BoundGenericType::get(ctx.getKeyPathDecl(), Type(),
-                                   {root, value});
-    } else if (objectTy->isPartialKeyPath() &&
-               kind == AllowedBindingKind::Subtypes) {
-      auto rootTy = objectTy->castTo<BoundGenericType>()->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 = CS.getKeyPathValueType(keyPath);
-
-      type = BoundGenericType::get(ctx.getKeyPathDecl(), Type(),
-                                   {rootTy, valueTy});
-    }
   }
 
   if (auto *locator = TypeVar->getImpl().getLocator()) {

lib/Sema/CSSimplify.cpp

@@ -11608,19 +11608,52 @@ bool ConstraintSystem::resolveKeyPath(TypeVariableType *typeVar,
                                       ConstraintLocatorBuilder locator) {
   auto *keyPathLocator = typeVar->getImpl().getLocator();
   auto *keyPath = castToExpr<KeyPathExpr>(keyPathLocator->getAnchor());
+
   if (keyPath->hasSingleInvalidComponent()) {
     assignFixedType(typeVar, contextualType);
     return true;
   }
-  if (auto *BGT = contextualType->getAs<BoundGenericType>()) {
-    auto args = BGT->getGenericArgs();
-    if (isKnownKeyPathType(contextualType) && args.size() >= 1) {
-      auto root = BGT->getGenericArgs()[0];
 
-      auto *keyPathValueTV = getKeyPathValueType(keyPath);
-      contextualType = BoundGenericType::get(
-          args.size() == 1 ? getASTContext().getKeyPathDecl() : BGT->getDecl(),
-          /*parent=*/Type(), {root, keyPathValueTV});
+  auto objectTy = contextualType->lookThroughAllOptionalTypes();
+  {
+    auto &ctx = getASTContext();
+    // `AnyKeyPath` and `PartialKeyPath` represent type-erased versions of
+    // `KeyPath<T, V>`.
+    //
+    // In situations where `AnyKeyPath` or `PartialKeyPath` cannot be used
+    // directly i.e. passing an argument to a parameter represented by a
+    // `AnyKeyPath` or `PartialKeyPath`, let's attempt a `KeyPath` binding which
+    // would then be converted to a `AnyKeyPath` or `PartialKeyPath` since there
+    // is a subtype relationship between them.
+    if (objectTy->isAnyKeyPath()) {
+      auto root = getKeyPathRootType(keyPath);
+      auto value = getKeyPathValueType(keyPath);
+
+      contextualType =
+          BoundGenericType::get(ctx.getKeyPathDecl(), Type(), {root, value});
+    } else if (objectTy->isPartialKeyPath()) {
+      auto rootTy = objectTy->castTo<BoundGenericType>()->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 = getKeyPathValueType(keyPath);
+
+      contextualType = BoundGenericType::get(ctx.getKeyPathDecl(), Type(),
+                                             {rootTy, valueTy});
+    } else if (isKnownKeyPathType(objectTy)) {
+      auto *keyPathTy = objectTy->castTo<BoundGenericType>();
+      auto args = keyPathTy->getGenericArgs();
+      assert(args.size() == 2);
+
+      auto root = args.front();
+      auto value = getKeyPathValueType(keyPath);
+
+      // Make sure that key path always gets a chance to infer its
+      // value type from the member chain.
+      if (!value->isEqual(args.back())) {
+        contextualType = BoundGenericType::get(
+            keyPathTy->getDecl(), keyPathTy->getParent(), {root, value});
+      }
     }
   }
 

validation-test/Sema/type_checker_crashers_fixed/rdar113760727.swift

@@ -0,0 +1,20 @@
+// RUN: %target-typecheck-verify-swift
+
+struct Value {
+  var data: String
+}
+
+func testAnyKeyPath<T>(field: AnyKeyPath, _: UnsafePointer<T>) {}
+func testPartialKeyPath<T>(field: PartialKeyPath<T>, _: UnsafePointer<T>) {}
+
+struct Test {
+  func test(ptr: UnsafePointer<Int64>) {
+    ptr.withMemoryRebound(to: Value.self, capacity: 2) { newPtr in
+      for idx in 0 ..< 2 {
+        var elt = newPtr[idx]
+        testAnyKeyPath(field: \Value.data, &elt)
+        testPartialKeyPath(field: \Value.data, &elt)
+      }
+    }
+  }
+}