[ConstraintSystem] Detect and fix invalid refs in dynamic key path me…

include/swift/AST/DiagnosticsSema.def

@@ -473,13 +473,15 @@ ERROR(expr_keypath_type_of_property,none,
 ERROR(expr_keypath_generic_type,none,
       "key path cannot refer to generic type %0", (DeclName))
 ERROR(expr_keypath_not_property,none,
-      "key path cannot refer to %0 %1", (DescriptiveDeclKind, DeclName))
+      "%select{key path|dynamic key path member lookup}2 cannot refer to %0 %1",
+      (DescriptiveDeclKind, DeclName, bool))
 ERROR(expr_keypath_mutating_getter,none,
-      "key path cannot refer to %0, which has a mutating getter",
-      (DeclName))
+      "%select{key path|dynamic key path member lookup}1 cannot refer to %0, "
+      "which has a mutating getter",
+      (DeclName, bool))
 ERROR(expr_keypath_static_member,none,
-      "key path cannot refer to static member %0",
-      (DeclName))
+      "%select{key path|dynamic key path member lookup}1 cannot refer to static member %0",
+      (DeclName, bool))
 ERROR(expr_keypath_empty,none,
       "empty key path does not refer to a property", ())
 ERROR(expr_unsupported_objc_key_path_component,none,

lib/Sema/CSDiagnostics.cpp

@@ -2532,17 +2532,19 @@ SourceLoc InvalidMemberRefInKeyPath::getLoc() const {
 }
 
 bool InvalidStaticMemberRefInKeyPath::diagnoseAsError() {
-  emitDiagnostic(getLoc(), diag::expr_keypath_static_member, getName());
+  emitDiagnostic(getLoc(), diag::expr_keypath_static_member, getName(),
+                 isForKeyPathDynamicMemberLookup());
   return true;
 }
 
 bool InvalidMemberWithMutatingGetterInKeyPath::diagnoseAsError() {
-  emitDiagnostic(getLoc(), diag::expr_keypath_mutating_getter, getName());
+  emitDiagnostic(getLoc(), diag::expr_keypath_mutating_getter, getName(),
+                 isForKeyPathDynamicMemberLookup());
   return true;
 }
 
 bool InvalidMethodRefInKeyPath::diagnoseAsError() {
   emitDiagnostic(getLoc(), diag::expr_keypath_not_property, getKind(),
-                 getName());
+                 getName(), isForKeyPathDynamicMemberLookup());
   return true;
 }

lib/Sema/CSDiagnostics.h

@@ -1039,7 +1039,8 @@ class InvalidMemberRefInKeyPath : public FailureDiagnostic {
                             ConstraintLocator *locator)
       : FailureDiagnostic(root, cs, locator), Member(member) {
     assert(member->hasName());
-    assert(locator->isForKeyPathComponent());
+    assert(locator->isForKeyPathComponent() ||
+           locator->isForKeyPathDynamicMemberLookup());
   }
 
   DescriptiveDeclKind getKind() const { return Member->getDescriptiveKind(); }
@@ -1051,6 +1052,10 @@ class InvalidMemberRefInKeyPath : public FailureDiagnostic {
 protected:
   /// Compute location of the failure for diagnostic.
   SourceLoc getLoc() const;
+
+  bool isForKeyPathDynamicMemberLookup() const {
+    return getLocator()->isForKeyPathDynamicMemberLookup();
+  }
 };
 
 /// Diagnose an attempt to reference a static member as a key path component

lib/Sema/CSSimplify.cpp

@@ -3710,11 +3710,6 @@ bool swift::hasDynamicMemberLookupAttribute(Type type) {
   return ::hasDynamicMemberLookupAttribute(type, DynamicMemberLookupCache);
 }
 
-static bool isKeyPathDynamicMemberLookup(ConstraintLocator *locator) {
-  auto path = locator ? locator->getPath() : None;
-  return !path.empty() && path.back().isKeyPathDynamicMember();
-}
-
 /// Given a ValueMember, UnresolvedValueMember, or TypeMember constraint,
 /// perform a lookup into the specified base type to find a candidate list.
 /// The list returned includes the viable candidates as well as the unviable
@@ -3753,7 +3748,7 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
   // is a single argument with `keypath:` label or `\.` syntax is used.
   if (memberName.isSimpleName() &&
       memberName.getBaseName().getKind() == DeclBaseName::Kind::Subscript &&
-      !isKeyPathDynamicMemberLookup(memberLocator)) {
+      !(memberLocator && memberLocator->isForKeyPathDynamicMemberLookup())) {
     if (baseTy->isAnyObject()) {
       result.addUnviable(
           OverloadChoice(baseTy, OverloadChoiceKind::KeyPathApplication),
@@ -4019,7 +4014,7 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
     // based dynamic member lookup. Since it's unknown upfront
     // what kind of declaration lookup is going to find, let's
     // double check here that given keypath is appropriate for it.
-    if (isKeyPathDynamicMemberLookup(memberLocator)) {
+    if (memberLocator && memberLocator->isForKeyPathDynamicMemberLookup()) {
       auto path = memberLocator->getPath();
       auto *keyPath = path.back().getKeyPath();
       if (auto *storage = dyn_cast<AbstractStorageDecl>(decl)) {
@@ -4416,10 +4411,17 @@ fixMemberRef(ConstraintSystem &cs, Type baseTy,
   // Not all of the choices handled here are going
   // to refer to a declaration.
   if (choice.isDecl()) {
-    if (auto *CD = dyn_cast<ConstructorDecl>(choice.getDecl())) {
+    auto *decl = choice.getDecl();
+
+    if (auto *CD = dyn_cast<ConstructorDecl>(decl)) {
       if (auto *fix = validateInitializerRef(cs, CD, locator))
         return fix;
     }
+
+    if (locator->isForKeyPathDynamicMemberLookup()) {
+      if (auto *fix = AllowInvalidRefInKeyPath::forRef(cs, decl, locator))
+        return fix;
+    }
   }
 
   if (reason) {

lib/Sema/ConstraintLocator.cpp

@@ -159,6 +159,11 @@ bool ConstraintLocator::isKeyPathSubscriptComponent() const {
   });
 }
 
+bool ConstraintLocator::isForKeyPathDynamicMemberLookup() const {
+  auto path = getPath();
+  return !path.empty() && path.back().isKeyPathDynamicMember();
+}
+
 bool ConstraintLocator::isForKeyPathComponent() const {
   return llvm::any_of(getPath(), [&](const LocatorPathElt &elt) {
     return elt.isKeyPathComponent();

lib/Sema/ConstraintLocator.h

@@ -560,6 +560,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// of the key path at some index.
   bool isKeyPathSubscriptComponent() const;
 
+  /// Determine whether this locator points to the member found
+  /// via key path dynamic member lookup.
+  bool isForKeyPathDynamicMemberLookup() const;
+
   /// Determine whether this locator points to one of the key path
   /// components.
   bool isForKeyPathComponent() const;

lib/Sema/TypeCheckExprObjC.cpp

@@ -395,7 +395,8 @@ Optional<Type> TypeChecker::checkObjCKeyPathExpr(DeclContext *dc,
 
     // Declarations that cannot be part of a key-path.
     diagnose(componentNameLoc, diag::expr_keypath_not_property,
-             found->getDescriptiveKind(), found->getFullName());
+             found->getDescriptiveKind(), found->getFullName(),
+             /*isForDynamicKeyPathMemberLookup=*/false);
     isInvalid = true;
     break;
   }

test/attr/attr_dynamic_member_lookup.swift

@@ -667,3 +667,24 @@ func test_chain_of_recursive_lookups(_ lens: Lens<Lens<Lens<Point>>>) {
   _ = \Lens<Lens<Point>>.x
   _ = \Lens<Lens<Point>>.obj.x
 }
+
+// KeyPath Dynamic Member Lookup can't refer to methods, mutating setters and static members
+// because of the KeyPath limitations
+func invalid_refs_through_dynamic_lookup() {
+  struct S {
+    static var foo: Int = 42
+    func bar() -> Q { return Q() }
+    static func baz(_: String) -> Int { return 0 }
+  }
+
+  struct Q {
+    var faz: Int = 0
+  }
+
+  func test(_ lens: A<S>) {
+    _ = lens.foo           // expected-error {{dynamic key path member lookup cannot refer to static member 'foo'}}
+    _ = lens.bar()         // expected-error {{dynamic key path member lookup cannot refer to instance method 'bar()'}}
+    _ = lens.bar().faz + 1 // expected-error {{dynamic key path member lookup cannot refer to instance method 'bar()'}}
+    _ = lens.baz("hello")  // expected-error {{dynamic key path member lookup cannot refer to static method 'baz'}}
+  }
+}