[ConstraintLocator] Identify whether locator is a result of keypath dynamic member lookup or belongs to keypath subscript component

Author: xedin

lib/Sema/CSSimplify.cpp

@@ -983,21 +983,6 @@ ConstraintSystem::TypeMatchResult constraints::matchCallArguments(
     const auto &param = params[paramIdx];
     auto paramTy = param.getOldType();
 
-    // Each of the index parameters has to conform to Hashable
-    // to be viable for use as a keypath subscript component.
-    if (keyPathSubscriptComponent) {
-      auto *hashable =
-          cs.getASTContext().getProtocol(KnownProtocolKind::Hashable);
-      // Standard library might be broken.
-      if (!hashable)
-        return cs.getTypeMatchFailure(locator);
-
-      cs.addConstraint(
-          ConstraintKind::ConformsTo, paramTy, hashable->getDeclaredType(),
-          cs.getConstraintLocator(keyPathSubscriptComponent,
-                                  LocatorPathElt::getTupleElement(paramIdx)));
-    }
-
     if (param.isAutoClosure())
       paramTy = paramTy->castTo<FunctionType>()->getResult();
 
@@ -3159,11 +3144,13 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyConformsToConstraint(
     // If this is an implicit Hashable conformance check generated for each
     // index argument of the keypath subscript component, we could just treat
     // it as though it conforms.
-    if (auto *component = getAsKeyPathSubscriptComponent(*this, anchor, path)) {
+    auto *loc = getConstraintLocator(locator);
+    if (loc->isResultOfKeyPathDynamicMemberLookup() ||
+        loc->isKeyPathSubscriptComponent()) {
       if (protocol ==
           getASTContext().getProtocol(KnownProtocolKind::Hashable)) {
-        auto *fix = TreatKeyPathSubscriptIndexAsHashable::create(*this, type,
-                                                                 component);
+        auto *fix =
+            TreatKeyPathSubscriptIndexAsHashable::create(*this, type, loc);
         if (!recordFix(fix))
           return SolutionKind::Solved;
       }

lib/Sema/ConstraintLocator.cpp

@@ -101,6 +101,30 @@ bool ConstraintLocator::isSubscriptMemberRef() const {
   return path.back().getKind() == ConstraintLocator::SubscriptMember;
 }
 
+bool ConstraintLocator::isResultOfKeyPathDynamicMemberLookup() const {
+  return llvm::any_of(getPath(), [](const LocatorPathElt &elt) {
+    return elt.isKeyPathDynamicMember();
+  });
+}
+
+bool ConstraintLocator::isKeyPathSubscriptComponent() const {
+  auto *anchor = getAnchor();
+  auto *KPE = dyn_cast_or_null<KeyPathExpr>(anchor);
+  if (!KPE)
+    return false;
+
+  using ComponentKind = KeyPathExpr::Component::Kind;
+  return llvm::any_of(getPath(), [&](const LocatorPathElt &elt) {
+    if (!elt.isKeyPathComponent())
+      return false;
+
+    auto index = elt.getValue();
+    auto &component = KPE->getComponents()[index];
+    return component.getKind() == ComponentKind::Subscript ||
+           component.getKind() == ComponentKind::UnresolvedSubscript;
+  });
+}
+
 void ConstraintLocator::dump(SourceManager *sm) {
   dump(sm, llvm::errs());
   llvm::errs() << "\n";

lib/Sema/ConstraintLocator.h

@@ -522,6 +522,14 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// e.g. `foo[0]` or `\Foo.[0]`
   bool isSubscriptMemberRef() const;
 
+  /// Determine whether given locator points to the choice picked as
+  /// as result of the key path dynamic member lookup operation.
+  bool isResultOfKeyPathDynamicMemberLookup() const;
+
+  /// Determine whether given locator points to a subscript component
+  /// of the key path at some index.
+  bool isKeyPathSubscriptComponent() const;
+
   /// Produce a profile of this locator, for use in a folding set.
   static void Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
                       ArrayRef<PathElement> path);

lib/Sema/ConstraintSystem.cpp

@@ -1695,36 +1695,6 @@ isInvalidPartialApplication(ConstraintSystem &cs, const ValueDecl *member,
   return {true, level};
 }
 
-static bool isResultOfKeyPathDynamicMemberLookup(ConstraintLocator *locator) {
-  auto path = locator->getPath();
-  return !path.empty() && path.back().isKeyPathDynamicMember();
-}
-
-/// If given anchor + path represents a key path with subscript
-/// component at some index.
-static bool isKeyPathSubscriptComponent(ConstraintLocator *locator) {
-  auto *anchor = locator->getAnchor();
-  auto *KPE = dyn_cast_or_null<KeyPathExpr>(anchor);
-  if (!KPE)
-    return false;
-
-  using ComponentKind = KeyPathExpr::Component::Kind;
-  auto path = locator->getPath();
-  while (!path.empty()) {
-    auto &last = path.back();
-    if (!last.isKeyPathComponent()) {
-      path = path.drop_back();
-      continue;
-    }
-
-    auto index = last.getValue();
-    auto &component = KPE->getComponents()[index];
-    return component.getKind() == ComponentKind::Subscript ||
-           component.getKind() == ComponentKind::UnresolvedSubscript;
-  }
-  return false;
-}
-
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                        Type boundType,
                                        OverloadChoice choice,
@@ -2122,8 +2092,8 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     }
 
     if (auto *SD = dyn_cast<SubscriptDecl>(decl)) {
-      if (isResultOfKeyPathDynamicMemberLookup(locator) ||
-          isKeyPathSubscriptComponent(locator)) {
+      if (locator->isResultOfKeyPathDynamicMemberLookup() ||
+          locator->isKeyPathSubscriptComponent()) {
         // Subscript type has a format of (Self[.Type) -> (Arg...) -> Result
         auto declTy = openedFullType->castTo<FunctionType>();
         auto subscriptTy = declTy->getResult()->castTo<FunctionType>();