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

Referencing (instance or static) methods in the key path is not
currently allowed, solver should be responsible for early detection
and diagnosis of both standalone e.g. `\.foo` and chained
e.g. `\.foo.bar` (where foo is a method) references in key path
components.

```swift
struct S {
  func foo() -> Int { return 42 }
}

let _: KeyPath<S, Int> = \.foo
```

Resolves: rdar://problem/49413561

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())

lib/Sema/CSFix.cpp

@@ -436,7 +436,36 @@ bool AllowInvalidRefInKeyPath::diagnose(Expr *root, bool asNote) const {
         root, getConstraintSystem(), Member, getLocator());
     return failure.diagnose(asNote);
   }
+
+  case RefKind::Method: {
+    return false;
+  }
+  }
+}
+
+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;