[CS] Add locator last element casting members

These members provide a convenient way of casting the last element of
a locator to a given path element type.

Author: hamishknight

lib/Sema/CSDiag.cpp

@@ -870,11 +870,9 @@ diagnoseUnresolvedDotExprTypeRequirementFailure(ConstraintSystem &cs,
   if (!locator)
     return false;
 
-  auto path = locator->getPath();
-  if (path.empty())
-    return false;
 
-  auto reqElt = path.back().getAs<LocatorPathElt::TypeParameterRequirement>();
+  auto reqElt =
+      locator->getLastElementAs<LocatorPathElt::TypeParameterRequirement>();
   if (!reqElt)
     return false;
 

lib/Sema/CSDiagnostics.cpp

@@ -330,13 +330,11 @@ const Requirement &RequirementFailure::getRequirement() const {
 ProtocolConformance *RequirementFailure::getConformanceForConditionalReq(
     ConstraintLocator *locator) {
   auto &cs = getConstraintSystem();
-  auto path = locator->getPath();
-  assert(!path.empty());
-  assert(path.back().is<LocatorPathElt::AnyRequirement>());
-
-  if (!path.back().isConditionalRequirement())
+  auto reqElt = locator->castLastElementTo<LocatorPathElt::AnyRequirement>();
+  if (!reqElt.isConditionalRequirement())
     return nullptr;
 
+  auto path = locator->getPath();
   auto *typeReqLoc = getConstraintLocator(getRawAnchor(), path.drop_back());
 
   auto result = llvm::find_if(
@@ -780,14 +778,9 @@ bool NoEscapeFuncToTypeConversionFailure::diagnoseAsError() {
     return true;
   }
 
-  GenericTypeParamType *paramTy = nullptr;
-
-  auto path = getLocator()->getPath();
-  if (!path.empty())
-    if (auto gpElt = path.back().getAs<LocatorPathElt::GenericParameter>())
-      paramTy = gpElt->getType();
-
-  if (paramTy) {
+  auto *loc = getLocator();
+  if (auto gpElt = loc->getLastElementAs<LocatorPathElt::GenericParameter>()) {
+    auto *paramTy = gpElt->getType();
     emitDiagnostic(anchor->getLoc(), diag::converting_noescape_to_type,
                   paramTy);
   } else {
@@ -1254,8 +1247,8 @@ bool RValueTreatedAsLValueFailure::diagnoseAsError() {
       auto argTuple = dyn_cast<TupleExpr>(argExpr);
       diagExpr = argTuple->getElement(0);
     } else if (getLocator()->getPath().size() > 0) {
-      auto lastPathElement = getLocator()->getPath().back();
-      auto argElt = lastPathElement.castTo<LocatorPathElt::ApplyArgToParam>();
+      auto argElt =
+          getLocator()->castLastElementTo<LocatorPathElt::ApplyArgToParam>();
 
       subElementDiagID = diag::cannot_pass_rvalue_inout_subelement;
       rvalueDiagID = diag::cannot_pass_rvalue_inout;
@@ -2406,10 +2399,8 @@ bool TupleContextualFailure::diagnoseAsError() {
 }
 
 bool AutoClosureForwardingFailure::diagnoseAsError() {
-  auto path = getLocator()->getPath();
-  assert(!path.empty());
-
-  auto last = path.back().castTo<LocatorPathElt::ApplyArgToParam>();
+  auto *loc = getLocator();
+  auto last = loc->castLastElementTo<LocatorPathElt::ApplyArgToParam>();
 
   // We need a raw anchor here because `getAnchor()` is simplified
   // to the argument expression.
@@ -3684,10 +3675,7 @@ bool CollectionElementContextualFailure::diagnoseAsError() {
   }
 
   if (isa<DictionaryExpr>(getRawAnchor())) {
-    auto path = locator->getPath();
-    assert(!path.empty());
-    auto eltLoc = path.back().castTo<LocatorPathElt::TupleElement>();
-
+    auto eltLoc = locator->castLastElementTo<LocatorPathElt::TupleElement>();
     switch (eltLoc.getIndex()) {
     case 0: // key
       diagnostic.emplace(emitDiagnostic(anchor->getLoc(),

lib/Sema/CSDiagnostics.h

@@ -237,11 +237,8 @@ class RequirementFailure : public FailureDiagnostic {
     assert(isConditional() || Signature);
     assert(AffectedDecl);
 
-    auto path = locator->getPath();
-    assert(!path.empty());
-
-    auto last = path.back().castTo<LocatorPathElt::AnyRequirement>();
-    assert(last.getRequirementKind() == kind);
+    auto reqElt = locator->castLastElementTo<LocatorPathElt::AnyRequirement>();
+    assert(reqElt.getRequirementKind() == kind);
 
     // It's possible sometimes not to have no base expression.
     if (!expr)
@@ -252,8 +249,8 @@ class RequirementFailure : public FailureDiagnostic {
   }
 
   unsigned getRequirementIndex() const {
-    auto path = getLocator()->getPath();
-    auto reqElt = path.back().castTo<LocatorPathElt::AnyRequirement>();
+    auto reqElt =
+        getLocator()->castLastElementTo<LocatorPathElt::AnyRequirement>();
     return reqElt.getIndex();
   }
 

lib/Sema/CSSimplify.cpp

@@ -4518,32 +4518,33 @@ 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 (memberLocator && memberLocator->isForKeyPathDynamicMemberLookup()) {
-      auto path = memberLocator->getPath();
-      auto kpElt = path.back().castTo<LocatorPathElt::KeyPathDynamicMember>();
-      auto *keyPath = kpElt.getKeyPathDecl();
-      if (auto *storage = dyn_cast<AbstractStorageDecl>(decl)) {
-        // If this is an attempt to access read-only member via
-        // writable key path, let's fail this choice early.
-        auto &ctx = getASTContext();
-        if (isReadOnlyKeyPathComponent(storage) &&
-            (keyPath == ctx.getWritableKeyPathDecl() ||
-             keyPath == ctx.getReferenceWritableKeyPathDecl())) {
-          result.addUnviable(
-              candidate,
-              MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember);
-          return;
-        }
+    if (memberLocator) {
+      using KPDynamicMemberElt = LocatorPathElt::KeyPathDynamicMember;
+      if (auto kpElt = memberLocator->getLastElementAs<KPDynamicMemberElt>()) {
+        auto *keyPath = kpElt->getKeyPathDecl();
+        if (auto *storage = dyn_cast<AbstractStorageDecl>(decl)) {
+          // If this is an attempt to access read-only member via
+          // writable key path, let's fail this choice early.
+          auto &ctx = getASTContext();
+          if (isReadOnlyKeyPathComponent(storage) &&
+              (keyPath == ctx.getWritableKeyPathDecl() ||
+               keyPath == ctx.getReferenceWritableKeyPathDecl())) {
+            result.addUnviable(
+                candidate,
+                MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember);
+            return;
+          }
 
-        // A nonmutating setter indicates a reference-writable base,
-        // on the other hand if setter is mutating there is no point
-        // of attempting `ReferenceWritableKeyPath` overload.
-        if (storage->isSetterMutating() &&
-            keyPath == ctx.getReferenceWritableKeyPathDecl()) {
-          result.addUnviable(
-              candidate,
-              MemberLookupResult::UR_ReferenceWritableKeyPathOnMutatingMember);
-          return;
+          // A nonmutating setter indicates a reference-writable base,
+          // on the other hand if setter is mutating there is no point
+          // of attempting `ReferenceWritableKeyPath` overload.
+          if (storage->isSetterMutating() &&
+              keyPath == ctx.getReferenceWritableKeyPathDecl()) {
+            result.addUnviable(candidate,
+                               MemberLookupResult::
+                                   UR_ReferenceWritableKeyPathOnMutatingMember);
+            return;
+          }
         }
       }
     }

lib/Sema/ConstraintLocator.cpp

@@ -195,9 +195,7 @@ bool ConstraintLocator::isForContextualType() const {
 }
 
 GenericTypeParamType *ConstraintLocator::getGenericParameter() const {
-  auto path = getPath();
-  assert(!path.empty());
-  return path.back().castTo<LocatorPathElt::GenericParameter>().getType();
+  return castLastElementTo<LocatorPathElt::GenericParameter>().getType();
 }
 
 void ConstraintLocator::dump(SourceManager *sm) {

lib/Sema/ConstraintLocator.h

@@ -510,6 +510,27 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// is of a given kind.
   bool isLastElement(ConstraintLocator::PathElementKind kind) const;
 
+  /// Attempts to cast the last path element of the locator to a specific
+  /// \c LocatorPathElt subclass, returning \c None if either unsuccessful or
+  /// the locator has no path elements.
+  template <class T>
+  Optional<T> getLastElementAs() const {
+    auto path = getPath();
+    if (path.empty())
+      return None;
+
+    return path.back().getAs<T>();
+  }
+
+  /// Casts the last path element of the locator to a specific \c LocatorPathElt
+  /// subclass, asserting that it has at least one element.
+  template <class T>
+  T castLastElementTo() const {
+    auto path = getPath();
+    assert(!path.empty() && "Expected at least one path element!");
+    return path.back().castTo<T>();
+  }
+
   /// If this locator points to generic parameter return its type.
   GenericTypeParamType *getGenericParameter() const;
 
@@ -578,7 +599,8 @@ using LocatorPathElt = ConstraintLocator::PathElement;
 
 // Disallow direct uses of isa/cast/dyn_cast on LocatorPathElt in favor of using
 // is/castTo/getAs. This allows us to work with Optional<T> rather than pointers
-// for getAs.
+// for getAs, and maintains consistency with ConstraintLocator's
+// isLastElement/castLastElementTo/getLastElementAs members.
 template <class X>
 inline bool
 isa(const LocatorPathElt &) = delete; // Use LocatorPathElt::is instead.