Merge pull request #24208 from xedin/diagnose-method-ref-in-keypath

[ConstraintSystem] Detect and fix use of method references in key path

Author: xedin

lib/Sema/CSApply.cpp

@@ -4573,20 +4573,13 @@ namespace {
                                   ConstraintLocator *locator) {
       if (overload.choice.isDecl()) {
         auto property = overload.choice.getDecl();
-
         // Key paths can only refer to properties currently.
-        if (!isa<VarDecl>(property)) {
-          cs.TC.diagnose(componentLoc, diag::expr_keypath_not_property,
-                         property->getDescriptiveKind(),
-                         property->getFullName());
-        } else {
-          // Key paths don't work with mutating-get properties.
-          auto varDecl = cast<VarDecl>(property);
-          assert(!varDecl->isGetterMutating());
-          // Key paths don't currently support static members.
-          // There is a fix which diagnoses such situation already.
-          assert(!varDecl->isStatic());
-        }
+        auto varDecl = cast<VarDecl>(property);
+        // Key paths don't work with mutating-get properties.
+        assert(!varDecl->isGetterMutating());
+        // Key paths don't currently support static members.
+        // There is a fix which diagnoses such situation already.
+        assert(!varDecl->isStatic());
 
         cs.TC.requestMemberLayout(property);
 

lib/Sema/CSDiagnostics.cpp

@@ -1516,6 +1516,13 @@ bool MissingCallFailure::diagnoseAsError() {
   if (auto *FVE = dyn_cast<ForceValueExpr>(baseExpr))
     baseExpr = FVE->getSubExpr();
 
+  // Calls are not yet supported by key path, but it
+  // is useful to record this fix to diagnose chaining
+  // where one of the key path components is a method
+  // reference.
+  if (isa<KeyPathExpr>(baseExpr))
+    return false;
+
   if (auto *DRE = dyn_cast<DeclRefExpr>(baseExpr)) {
     emitDiagnostic(baseExpr->getLoc(), diag::did_not_call_function,
                    DRE->getDecl()->getBaseName().getIdentifier())
@@ -2479,3 +2486,9 @@ bool InvalidMemberWithMutatingGetterInKeyPath::diagnoseAsError() {
   emitDiagnostic(getLoc(), diag::expr_keypath_mutating_getter, getName());
   return true;
 }
+
+bool InvalidMethodRefInKeyPath::diagnoseAsError() {
+  emitDiagnostic(getLoc(), diag::expr_keypath_not_property, getKind(),
+                 getName());
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -1042,6 +1042,8 @@ class InvalidMemberRefInKeyPath : public FailureDiagnostic {
     assert(locator->isForKeyPathComponent());
   }
 
+  DescriptiveDeclKind getKind() const { return Member->getDescriptiveKind(); }
+
   DeclName getName() const { return Member->getFullName(); }
 
   bool diagnoseAsError() override = 0;
@@ -1098,6 +1100,29 @@ class InvalidMemberWithMutatingGetterInKeyPath final
   bool diagnoseAsError() override;
 };
 
+/// Diagnose an attempt to reference a method as a key path component
+/// e.g.
+///
+/// ```swift
+/// struct S {
+///   func foo() -> Int { return 42 }
+///   static func bar() -> Int { return 0 }
+/// }
+///
+/// _ = \S.foo
+/// _ = \S.Type.bar
+/// ```
+class InvalidMethodRefInKeyPath final : public InvalidMemberRefInKeyPath {
+public:
+  InvalidMethodRefInKeyPath(Expr *root, ConstraintSystem &cs, ValueDecl *method,
+                            ConstraintLocator *locator)
+      : InvalidMemberRefInKeyPath(root, cs, method, locator) {
+    assert(isa<FuncDecl>(method));
+  }
+
+  bool diagnoseAsError() override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -436,7 +436,38 @@ bool AllowInvalidRefInKeyPath::diagnose(Expr *root, bool asNote) const {
         root, getConstraintSystem(), Member, getLocator());
     return failure.diagnose(asNote);
   }
+
+  case RefKind::Method: {
+    InvalidMethodRefInKeyPath failure(root, getConstraintSystem(), Member,
+                                      getLocator());
+    return failure.diagnose(asNote);
+  }
+  }
+}
+
+AllowInvalidRefInKeyPath *
+AllowInvalidRefInKeyPath::forRef(ConstraintSystem &cs, ValueDecl *member,
+                                 ConstraintLocator *locator) {
+  // Referencing (instance or static) methods in key path is
+  // not currently allowed.
+  if (isa<FuncDecl>(member))
+    return AllowInvalidRefInKeyPath::create(cs, RefKind::Method, member,
+                                            locator);
+
+  // Referencing static members in key path is not currently allowed.
+  if (member->isStatic())
+    return AllowInvalidRefInKeyPath::create(cs, RefKind::StaticMember, member,
+                                            locator);
+
+  if (auto *storage = dyn_cast<AbstractStorageDecl>(member)) {
+    // Referencing members with mutating getters in key path is not
+    // currently allowed.
+    if (storage->isGetterMutating())
+      return AllowInvalidRefInKeyPath::create(cs, RefKind::MutatingGetter,
+                                              member, locator);
   }
+
+  return nullptr;
 }
 
 AllowInvalidRefInKeyPath *

lib/Sema/CSFix.h

@@ -791,6 +791,9 @@ class AllowInvalidRefInKeyPath final : public ConstraintFix {
     // Allow a reference to a declaration with mutating getter as
     // a key path component.
     MutatingGetter,
+    // Allow a reference to a method (instance or static) as
+    // a key path component.
+    Method,
   } Kind;
 
   ValueDecl *Member;
@@ -808,22 +811,16 @@ class AllowInvalidRefInKeyPath final : public ConstraintFix {
     case RefKind::MutatingGetter:
       return "allow reference to a member with mutating getter as a key "
              "path component";
+    case RefKind::Method:
+      return "allow reference to a method as a key path component";
     }
   }
 
   bool diagnose(Expr *root, bool asNote = false) const override;
 
-  static AllowInvalidRefInKeyPath *forStaticMember(ConstraintSystem &cs,
-                                                   ValueDecl *member,
-                                                   ConstraintLocator *locator) {
-    return create(cs, RefKind::StaticMember, member, locator);
-  }
-
+  /// Determine whether give reference requires a fix and produce one.
   static AllowInvalidRefInKeyPath *
-  forMutatingGetter(ConstraintSystem &cs, ValueDecl *member,
-                    ConstraintLocator *locator) {
-    return create(cs, RefKind::MutatingGetter, member, locator);
-  }
+  forRef(ConstraintSystem &cs, ValueDecl *member, ConstraintLocator *locator);
 
 private:
   static AllowInvalidRefInKeyPath *create(ConstraintSystem &cs, RefKind kind,

lib/Sema/CSSimplify.cpp

@@ -1975,15 +1975,29 @@ static ConstraintFix *fixRequirementFailure(ConstraintSystem &cs, Type type1,
   }
 }
 
-static void
+/// Attempt to repair typing failures and record fixes if needed.
+/// \return true if at least some of the failures has been repaired
+/// successfully, which allows type matcher to continue.
+static bool
 repairFailures(ConstraintSystem &cs, Type lhs, Type rhs,
                SmallVectorImpl<RestrictionOrFix> &conversionsOrFixes,
                ConstraintLocatorBuilder locator) {
   SmallVector<LocatorPathElt, 4> path;
   auto *anchor = locator.getLocatorParts(path);
 
-  if (path.empty())
-    return;
+  if (path.empty()) {
+    // If method reference forms a value type of the key path,
+    // there is going to be a constraint to match result of the
+    // member lookup to the generic parameter `V` of *KeyPath<R, V>
+    // type associated with key path expression, which we need to
+    // fix-up here.
+    if (anchor && isa<KeyPathExpr>(anchor)) {
+      auto *fnType = lhs->getAs<FunctionType>();
+      if (fnType && fnType->getResult()->isEqual(rhs))
+        return true;
+    }
+    return false;
+  }
 
   auto &elt = path.back();
   switch (elt.getKind()) {
@@ -2007,7 +2021,7 @@ repairFailures(ConstraintSystem &cs, Type lhs, Type rhs,
     if (path.empty() ||
         !(path.back().getKind() == ConstraintLocator::ApplyArgToParam ||
           path.back().getKind() == ConstraintLocator::ContextualType))
-      return;
+      return false;
 
     auto arg = llvm::find_if(cs.getTypeVariables(),
                              [&argLoc](const TypeVariableType *typeVar) {
@@ -2058,13 +2072,14 @@ repairFailures(ConstraintSystem &cs, Type lhs, Type rhs,
                                              cs.getConstraintLocator(locator));
       conversionsOrFixes.push_back(fix);
     }
-
     break;
   }
 
   default:
-    return;
+    break;
   }
+
+  return !conversionsOrFixes.empty();
 }
 
 ConstraintSystem::TypeMatchResult
@@ -2873,8 +2888,12 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
   }
 
   // Attempt to repair any failures identifiable at this point.
-  if (attemptFixes)
-    repairFailures(*this, type1, type2, conversionsOrFixes, locator);
+  if (attemptFixes) {
+    if (repairFailures(*this, type1, type2, conversionsOrFixes, locator)) {
+      if (conversionsOrFixes.empty())
+        return getTypeMatchSuccess();
+    }
+  }
 
   if (conversionsOrFixes.empty())
     return getTypeMatchFailure(locator);
@@ -4368,7 +4387,7 @@ fixMemberRef(ConstraintSystem &cs, Type baseTy,
       break;
     case MemberLookupResult::UR_KeyPathWithAnyObjectRootType:
       return AllowAnyObjectKeyPathRoot::create(cs, locator);
-    }
+   }
   }
 
   return nullptr;
@@ -5124,31 +5143,25 @@ ConstraintSystem::simplifyKeyPathConstraint(Type keyPathTy,
       }
 
       auto storage = dyn_cast<AbstractStorageDecl>(choices[i].getDecl());
-      if (!storage) {
-        return SolutionKind::Error;
-      }
 
-      // Referencing static members in key path is not currently allowed.
-      if (storage->isStatic() || storage->isGetterMutating()) {
-        if (!shouldAttemptFixes())
-          return SolutionKind::Error;
+      auto *componentLoc = getConstraintLocator(
+          locator.withPathElement(LocatorPathElt::getKeyPathComponent(i)));
 
-        auto *componentLoc = getConstraintLocator(
-            locator.withPathElement(LocatorPathElt::getKeyPathComponent(i)));
+      if (auto *fix = AllowInvalidRefInKeyPath::forRef(
+              *this, choices[i].getDecl(), componentLoc)) {
+        if (!shouldAttemptFixes() || recordFix(fix))
+          return SolutionKind::Error;
 
-        ConstraintFix *fix = nullptr;
-        if (storage->isStatic()) {
-          fix = AllowInvalidRefInKeyPath::forStaticMember(
-              *this, choices[i].getDecl(), componentLoc);
-        } else {
-          fix = AllowInvalidRefInKeyPath::forMutatingGetter(
-              *this, choices[i].getDecl(), componentLoc);
+        // If this was a method reference let's mark it as read-only.
+        if (!storage) {
+          capability = ReadOnly;
+          continue;
         }
-
-        if (recordFix(fix))
-          return SolutionKind::Error;
       }
 
+      if (!storage)
+        return SolutionKind::Error;
+
       if (isReadOnlyKeyPathComponent(storage)) {
         capability = ReadOnly;
         continue;

test/expr/unary/keypath/keypath.swift

@@ -776,6 +776,31 @@ func test_keypath_with_mutating_getter() {
   }
 }
 
+func test_keypath_with_method_refs() {
+  struct S {
+    func foo() -> Int { return 42 }
+    static func bar() -> Int { return 0 }
+  }
+
+  let _: KeyPath<S, Int> = \.foo // expected-error {{key path cannot refer to instance method 'foo()'}}
+  let _: KeyPath<S, Int> = \.bar // expected-error {{key path cannot refer to static member 'bar()'}}
+  let _ = \S.Type.bar // expected-error {{key path cannot refer to static method 'bar()'}}
+
+  struct A {
+    func foo() -> B { return B() }
+    static func faz() -> B { return B() }
+  }
+
+  struct B {
+    var bar: Int = 42
+  }
+
+  let _: KeyPath<A, Int> = \.foo.bar // expected-error {{key path cannot refer to instance method 'foo()'}}
+  let _: KeyPath<A, Int> = \.faz.bar // expected-error {{key path cannot refer to static member 'faz()'}}
+  let _ = \A.foo.bar // expected-error {{key path cannot refer to instance method 'foo()'}}
+  let _ = \A.Type.faz.bar // expected-error {{key path cannot refer to static method 'faz()'}}
+}
+
 func testSyntaxErrors() { // expected-note{{}}
   _ = \.  ; // expected-error{{expected member name following '.'}}
   _ = \.a ;