[ConstraintLocator] Add information about expected value type to key path locator

Record a "value" type associated with this location to be able to
reference during inference.

Author: xedin

include/swift/Sema/ConstraintLocator.h

@@ -910,6 +910,22 @@ class LocatorPathElt::ContextualType final : public StoredIntegerElement<1> {
   }
 };
 
+class LocatorPathElt::KeyPathType final
+    : public StoredPointerElement<TypeBase> {
+public:
+  KeyPathType(Type valueType)
+      : StoredPointerElement(PathElementKind::KeyPathType,
+                             valueType.getPointer()) {
+    assert(valueType);
+  }
+
+  Type getValueType() const { return getStoredPointer(); }
+
+  static bool classof(const LocatorPathElt *elt) {
+    return elt->getKind() == PathElementKind::KeyPathType;
+  }
+};
+
 namespace details {
   template <typename CustomPathElement>
   class PathElement {

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -114,7 +114,7 @@ CUSTOM_LOCATOR_PATH_ELT(KeyPathDynamicMember)
 SIMPLE_LOCATOR_PATH_ELT(KeyPathRoot)
 
 /// The type of the key path expression.
-SIMPLE_LOCATOR_PATH_ELT(KeyPathType)
+CUSTOM_LOCATOR_PATH_ELT(KeyPathType)
 
 /// The value of a key path.
 SIMPLE_LOCATOR_PATH_ELT(KeyPathValue)

lib/Sema/CSGen.cpp

@@ -3258,10 +3258,11 @@ namespace {
       // The result is a KeyPath from the root to the end component.
       // The type of key path depends on the overloads chosen for the key
       // path components.
-      auto typeLoc =
-          CS.getConstraintLocator(locator, ConstraintLocator::KeyPathType);
+      auto typeLoc = CS.getConstraintLocator(
+          locator, LocatorPathElt::KeyPathType(rvalueBase));
+
       Type kpTy = CS.createTypeVariable(typeLoc, TVO_CanBindToNoEscape |
-                                        TVO_CanBindToHole);
+                                                     TVO_CanBindToHole);
       CS.addKeyPathConstraint(kpTy, root, rvalueBase, componentTypeVars,
                               locator);
       return kpTy;

lib/Sema/ConstraintLocator.cpp

@@ -122,7 +122,7 @@ bool ConstraintLocator::isKeyPathType() const {
   // The format of locator should be `<keypath expr> -> key path type`
   if (!anchor || !isExpr<KeyPathExpr>(anchor) || path.size() != 1)
     return false;
-  return path.back().getKind() == ConstraintLocator::KeyPathType;
+  return path.back().is<LocatorPathElt::KeyPathType>();
 }
 
 bool ConstraintLocator::isKeyPathRoot() const {