[CS] Add locator find[First/Last] members

These members provide a convenient way of getting the first or last
occurrence of a given kind of element in a locator's path.

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -1545,15 +1545,6 @@ namespace {
           LocatorPathElt::KeyPathDynamicMember(keyPathTy->getAnyNominal()));
       auto overload = solution.getOverloadChoice(componentLoc);
 
-      auto getKeyPathComponentIndex =
-          [](ConstraintLocator *locator) -> unsigned {
-        for (const auto &elt : locator->getPath()) {
-          if (auto kpElt = elt.getAs<LocatorPathElt::KeyPathComponent>())
-            return kpElt->getIndex();
-        }
-        llvm_unreachable("no keypath component node");
-      };
-
       // Looks like there is a chain of implicit `subscript(dynamicMember:)`
       // calls necessary to resolve a member reference.
       if (overload.choice.getKind() ==
@@ -1578,8 +1569,9 @@ namespace {
       // of a keypath expression e.g. `\Lens<[Int]>.count` where
       // `count` is referenced using dynamic lookup.
       if (auto *KPE = dyn_cast<KeyPathExpr>(anchor)) {
-        auto componentIdx = getKeyPathComponentIndex(memberLoc);
-        auto &origComponent = KPE->getComponents()[componentIdx];
+        auto kpElt = memberLoc->findFirst<LocatorPathElt::KeyPathComponent>();
+        assert(kpElt && "no keypath component node");
+        auto &origComponent = KPE->getComponents()[kpElt->getIndex()];
 
         using ComponentKind = KeyPathExpr::Component::Kind;
         if (origComponent.getKind() == ComponentKind::UnresolvedProperty) {

lib/Sema/CSDiagnostics.cpp

@@ -197,21 +197,18 @@ FailureDiagnostic::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
   // have to look through other elements that are generated from an argument
   // conversion such as GenericArgument for an optional-to-optional conversion,
   // and OptionalPayload for a value-to-optional conversion.
-  auto applyArgEltIter =
-      std::find_if(path.rbegin(), path.rend(), [](LocatorPathElt elt) {
-        return elt.is<LocatorPathElt::ApplyArgToParam>();
-      });
-
-  if (applyArgEltIter == path.rend())
+  auto iter = path.rbegin();
+  auto applyArgElt = locator->findLast<LocatorPathElt::ApplyArgToParam>(iter);
+  if (!applyArgElt)
     return None;
 
-  assert(std::find_if(applyArgEltIter + 1, path.rend(), [](LocatorPathElt elt) {
-    return elt.getKind() == ConstraintLocator::ApplyArgToParam;
-  }) == path.rend() && "Multiple ApplyArgToParam components?");
+  auto nextIter = iter + 1;
+  assert(!locator->findLast<LocatorPathElt::ApplyArgToParam>(nextIter) &&
+         "Multiple ApplyArgToParam components?");
 
   // Form a new locator that ends at the apply-arg-to-param element, and
   // simplify it to get the full argument expression.
-  auto argPath = path.drop_back(applyArgEltIter - path.rbegin());
+  auto argPath = path.drop_back(iter - path.rbegin());
   auto *argLocator = cs.getConstraintLocator(
       anchor, argPath, ConstraintLocator::getSummaryFlagsForPath(argPath));
 
@@ -271,9 +268,8 @@ FailureDiagnostic::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
     fnInterfaceType = resolveInterfaceType(rawFnType);
   }
 
-  auto applyArgElt = applyArgEltIter->castTo<LocatorPathElt::ApplyArgToParam>();
-  auto argIdx = applyArgElt.getArgIdx();
-  auto paramIdx = applyArgElt.getParamIdx();
+  auto argIdx = applyArgElt->getArgIdx();
+  auto paramIdx = applyArgElt->getParamIdx();
 
   return FunctionArgApplyInfo(argExpr, argIdx, getType(argExpr), paramIdx,
                               fnInterfaceType, fnType, callee);
@@ -409,12 +405,8 @@ GenericSignature *RequirementFailure::getSignature(ConstraintLocator *locator) {
   if (isConditional())
     return Conformance->getGenericSignature();
 
-  auto path = locator->getPath();
-  for (auto iter = path.rbegin(); iter != path.rend(); ++iter) {
-    const auto &elt = *iter;
-    if (auto genericElt = elt.getAs<LocatorPathElt::OpenedGeneric>())
-      return genericElt->getSignature();
-  }
+  if (auto genericElt = locator->findLast<LocatorPathElt::OpenedGeneric>())
+    return genericElt->getSignature();
 
   llvm_unreachable("Type requirement failure should always have signature");
 }
@@ -3563,12 +3555,8 @@ bool KeyPathSubscriptIndexHashableFailure::diagnoseAsError() {
   auto loc = anchor->getLoc();
   if (locator->isKeyPathSubscriptComponent()) {
     auto *KPE = cast<KeyPathExpr>(anchor);
-    for (auto &elt : locator->getPath()) {
-      if (auto kpElt = elt.getAs<LocatorPathElt::KeyPathComponent>()) {
-        loc = KPE->getComponents()[kpElt->getIndex()].getLoc();
-        break;
-      }
-    }
+    if (auto kpElt = locator->findFirst<LocatorPathElt::KeyPathComponent>())
+      loc = KPE->getComponents()[kpElt->getIndex()].getLoc();
   }
 
   emitDiagnostic(loc, diag::expr_keypath_subscript_index_not_hashable,
@@ -3581,14 +3569,9 @@ SourceLoc InvalidMemberRefInKeyPath::getLoc() const {
 
   if (auto *KPE = dyn_cast<KeyPathExpr>(anchor)) {
     auto *locator = getLocator();
-    auto componentIter =
-        llvm::find_if(locator->getPath(), [](const LocatorPathElt &elt) {
-          return elt.isKeyPathComponent();
-        });
-
-    assert(componentIter != locator->getPath().end());
-    auto component = componentIter->castTo<LocatorPathElt::KeyPathComponent>();
-    return KPE->getComponents()[component.getIndex()].getLoc();
+    auto component = locator->findFirst<LocatorPathElt::KeyPathComponent>();
+    assert(component);
+    return KPE->getComponents()[component->getIndex()].getLoc();
   }
 
   return anchor->getLoc();

lib/Sema/ConstraintLocator.h

@@ -531,6 +531,63 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     return path.back().castTo<T>();
   }
 
+  using PathIterator = ArrayRef<PathElement>::iterator;
+  using PathReverseIterator = ArrayRef<PathElement>::reverse_iterator;
+
+  /// Attempts to find the first element in the locator's path that is a
+  /// specific \c LocatorPathElt subclass, returning \c None if no such element
+  /// exists.
+  ///
+  /// \param iter A reference to an iterator which will be used to iterate
+  /// over the locator's path.
+  template <class T>
+  Optional<T> findFirst(PathIterator &iter) const {
+    auto path = getPath();
+    auto end = path.end();
+    assert(iter >= path.begin() && iter <= end);
+
+    for (; iter != end; ++iter)
+      if (auto elt = iter->getAs<T>())
+        return elt;
+    return None;
+  }
+
+  /// Attempts to find the first element in the locator's path that is a
+  /// specific \c LocatorPathElt subclass, returning \c None if no such element
+  /// exists.
+  template <class T>
+  Optional<T> findFirst() const {
+    auto iter = getPath().begin();
+    return findFirst<T>(iter);
+  }
+
+  /// Attempts to find the last element in the locator's path that is a
+  /// specific \c LocatorPathElt subclass, returning \c None if no such element
+  /// exists.
+  ///
+  /// \param iter A reference to a reverse iterator which will be used to
+  /// iterate over the locator's path.
+  template <class T>
+  Optional<T> findLast(PathReverseIterator &iter) const {
+    auto path = getPath();
+    auto end = path.rend();
+    assert(iter >= path.rbegin() && iter <= end);
+
+    for (; iter != end; ++iter)
+      if (auto elt = iter->getAs<T>())
+        return elt;
+    return None;
+  }
+
+  /// Attempts to find the last element in the locator's path that is a
+  /// specific \c LocatorPathElt subclass, returning \c None if no such element
+  /// exists.
+  template <class T>
+  Optional<T> findLast() const {
+    auto iter = getPath().rbegin();
+    return findLast<T>(iter);
+  }
+
   /// If this locator points to generic parameter return its type.
   GenericTypeParamType *getGenericParameter() const;