[ConstraintSystem] Detect and fix invalid refs in dynamic key path member lookup

KeyPath dynamic member lookup is limited to what key path itself
could do, so let's detect and diagnose invalid references just
like we do for regular key path expressions.

Resolves: rdar://problem/50376224

Author: xedin

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'}}
+  }
+}