[ConstraintSystem] Type of key path expression should be a known KeyPath type

lib/Sema/CSBindings.cpp

@@ -276,6 +276,15 @@ ConstraintSystem::getPotentialBindingForRelationalConstraint(
   if (type->hasError())
     return None;
 
+  if (auto *locator = typeVar->getImpl().getLocator()) {
+    if (locator->isKeyPathType()) {
+      auto *BGT =
+          type->lookThroughAllOptionalTypes()->getAs<BoundGenericType>();
+      if (!BGT || !isKnownKeyPathDecl(getASTContext(), BGT->getDecl()))
+        return None;
+    }
+  }
+
   // If the source of the binding is 'OptionalObject' constraint
   // and type variable is on the left-hand side, that means
   // that it _has_ to be of optional type, since the right-hand

lib/Sema/CSGen.cpp

@@ -3084,7 +3084,9 @@ namespace {
       } else {
         // The type of key path depends on the overloads chosen for the key
         // path components.
-        kpTy = CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
+        auto typeLoc =
+            CS.getConstraintLocator(locator, ConstraintLocator::KeyPathType);
+        kpTy = CS.createTypeVariable(typeLoc, TVO_CanBindToNoEscape);
         CS.addKeyPathConstraint(kpTy, root, rvalueBase, locator);
       }
       return kpTy;

lib/Sema/ConstraintLocator.cpp

@@ -79,6 +79,7 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
     case DynamicLookupResult:
     case ContextualType:
     case SynthesizedArgument:
+    case KeyPathType:
     case KeyPathRoot:
     case KeyPathValue:
     case KeyPathComponentResult:
@@ -104,6 +105,15 @@ bool ConstraintLocator::isSubscriptMemberRef() const {
   return path.back().getKind() == ConstraintLocator::SubscriptMember;
 }
 
+bool ConstraintLocator::isKeyPathType() const {
+  auto *anchor = getAnchor();
+  auto path = getPath();
+  // The format of locator should be `<keypath expr> -> key path type`
+  if (!anchor || !isa<KeyPathExpr>(anchor) || path.size() != 1)
+    return false;
+  return path.back().getKind() == ConstraintLocator::KeyPathType;
+}
+
 bool ConstraintLocator::isKeyPathRoot() const {
   auto *anchor = getAnchor();
   auto path = getPath();
@@ -340,6 +350,10 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) {
       out << "key path dynamic member lookup";
       break;
 
+    case KeyPathType:
+      out << "key path type";
+      break;
+
     case KeyPathRoot:
       out << "key path root";
       break;

lib/Sema/ConstraintLocator.h

@@ -129,9 +129,11 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     SynthesizedArgument,
     /// The member looked up via keypath based dynamic lookup.
     KeyPathDynamicMember,
-    /// The root of a keypath
+    /// The type of the key path expression
+    KeyPathType,
+    /// The root of a key path
     KeyPathRoot,
-    /// The value of a keypath
+    /// The value of a key path
     KeyPathValue,
     /// The result type of a key path component. Not used for subscripts.
     KeyPathComponentResult,
@@ -165,6 +167,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ImplicitlyUnwrappedDisjunctionChoice:
     case DynamicLookupResult:
     case ContextualType:
+    case KeyPathType:
     case KeyPathRoot:
     case KeyPathValue:
     case KeyPathComponentResult:
@@ -234,6 +237,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ContextualType:
     case SynthesizedArgument:
     case KeyPathDynamicMember:
+    case KeyPathType:
     case KeyPathRoot:
     case KeyPathValue:
     case KeyPathComponentResult:
@@ -534,6 +538,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// e.g. `foo[0]` or `\Foo.[0]`
   bool isSubscriptMemberRef() const;
 
+  /// Determine whether give locator points to the type of the
+  /// key path expression.
+  bool isKeyPathType() const;
+
   /// Determine whether given locator points to the keypath root
   bool isKeyPathRoot() const;
 

lib/Sema/ConstraintSystem.cpp

@@ -1913,12 +1913,8 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
       refType = fnType->getResult();
 
       auto *keyPathDecl = keyPathTy->getAnyNominal();
-      assert(
-          keyPathDecl &&
-          (keyPathDecl == getASTContext().getKeyPathDecl() ||
-           keyPathDecl == getASTContext().getWritableKeyPathDecl() ||
-           keyPathDecl == getASTContext().getReferenceWritableKeyPathDecl()) &&
-          "parameter is supposed to be a keypath");
+      assert(isKnownKeyPathDecl(getASTContext(), keyPathDecl) &&
+             "parameter is supposed to be a keypath");
 
       auto *keyPathLoc = getConstraintLocator(
           locator, LocatorPathElt::getKeyPathDynamicMember(keyPathDecl));
@@ -2712,3 +2708,9 @@ void ConstraintSystem::generateConstraints(
     recordChoice(constraints, index, choices[index]);
   }
 }
+
+bool constraints::isKnownKeyPathDecl(ASTContext &ctx, ValueDecl *decl) {
+  return decl == ctx.getKeyPathDecl() || decl == ctx.getWritableKeyPathDecl() ||
+         decl == ctx.getReferenceWritableKeyPathDecl() ||
+         decl == ctx.getPartialKeyPathDecl() || decl == ctx.getAnyKeyPathDecl();
+}

lib/Sema/ConstraintSystem.h

@@ -4010,6 +4010,10 @@ class DisjunctionChoiceProducer : public BindingProducer<DisjunctionChoice> {
                              IsExplicitConversion, isBeginningOfPartition);
   }
 };
+
+/// Determine whether given declaration is one for a key path
+/// `{Writable, ReferenceWritable}KeyPath`.
+bool isKnownKeyPathDecl(ASTContext &ctx, ValueDecl *decl);
 } // end namespace constraints
 
 template<typename ...Args>

test/expr/unary/keypath/keypath.swift

@@ -801,6 +801,31 @@ func test_keypath_with_method_refs() {
   let _ = \A.Type.faz.bar // expected-error {{key path cannot refer to static method 'faz()'}}
 }
 
+// SR-10467 - Argument type 'KeyPath<String, Int>' does not conform to expected type 'Any'
+func test_keypath_in_any_context() {
+  func foo(_: Any) {}
+  _ = foo(\String.count) // Ok
+}
+
+protocol PWithTypeAlias {
+  typealias Key = WritableKeyPath<Self, Int?>
+  static var fooKey: Key? { get }
+  static var barKey: Key! { get }
+  static var fazKey: Key?? { get }
+  static var bazKey: Key?! { get }
+}
+
+func test_keypath_inference_with_optionals() {
+  final class S : PWithTypeAlias {
+    static var fooKey: Key? { return \.foo }
+    static var barKey: Key! { return \.foo }
+    static var fazKey: Key?? { return \.foo }
+    static var bazKey: Key?! { return \.foo }
+
+    var foo: Int? = nil
+  }
+}
+
 func testSyntaxErrors() { // expected-note{{}}
   _ = \.  ; // expected-error{{expected member name following '.'}}
   _ = \.a ;