[ConstraintSystem] Update key path subscript index type variable to use new locator element

Author: xedin

include/swift/Sema/ConstraintSystem.h

@@ -499,6 +499,10 @@ class TypeVariableType::Implementation {
   /// expression.
   bool isKeyPathValue() const;
 
+  /// Determine whether this type variable represents an index parameter of
+  /// a special `subscript(keyPath:)` subscript.
+  bool isKeyPathSubscriptIndex() const;
+
   /// Determine whether this type variable represents a subscript result type.
   bool isSubscriptResultType() const;
 

lib/Sema/ConstraintSystem.cpp

@@ -3687,7 +3687,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     // The element type is T or @lvalue T based on the key path subtype and
     // the mutability of the base.
     auto keyPathIndexTy = createTypeVariable(
-        getConstraintLocator(locator, ConstraintLocator::FunctionArgument),
+        getConstraintLocator(locator, ConstraintLocator::KeyPathSubscriptIndex),
         TVO_CanBindToInOut);
     auto elementTy = createTypeVariable(
             getConstraintLocator(locator, ConstraintLocator::FunctionArgument),

lib/Sema/TypeCheckConstraints.cpp

@@ -152,6 +152,11 @@ bool TypeVariableType::Implementation::isKeyPathValue() const {
   return locator && locator->isKeyPathValue();
 }
 
+bool TypeVariableType::Implementation::isKeyPathSubscriptIndex() const {
+  return locator &&
+         locator->isLastElement<LocatorPathElt::KeyPathSubscriptIndex>();
+}
+
 bool TypeVariableType::Implementation::isSubscriptResultType() const {
   if (!(locator && locator->getAnchor()))
     return false;