[ConstraintSystem] Add a special locator for type of the key path expression

Author: xedin

lib/Sema/CSGen.cpp

@@ -3084,7 +3084,9 @@ namespace {
       } else {
         // The type of key path depends on the overloads chosen for the key
         // path components.
-        kpTy = CS.createTypeVariable(locator, TVO_CanBindToNoEscape);
+        auto typeLoc =
+            CS.getConstraintLocator(locator, ConstraintLocator::KeyPathType);
+        kpTy = CS.createTypeVariable(typeLoc, TVO_CanBindToNoEscape);
         CS.addKeyPathConstraint(kpTy, root, rvalueBase, locator);
       }
       return kpTy;

lib/Sema/ConstraintLocator.cpp

@@ -79,6 +79,7 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
     case DynamicLookupResult:
     case ContextualType:
     case SynthesizedArgument:
+    case KeyPathType:
     case KeyPathRoot:
     case KeyPathValue:
     case KeyPathComponentResult:
@@ -104,6 +105,15 @@ bool ConstraintLocator::isSubscriptMemberRef() const {
   return path.back().getKind() == ConstraintLocator::SubscriptMember;
 }
 
+bool ConstraintLocator::isKeyPathType() const {
+  auto *anchor = getAnchor();
+  auto path = getPath();
+  // The format of locator should be `<keypath expr> -> key path type`
+  if (!anchor || !isa<KeyPathExpr>(anchor) || path.size() != 1)
+    return false;
+  return path.back().getKind() == ConstraintLocator::KeyPathType;
+}
+
 bool ConstraintLocator::isKeyPathRoot() const {
   auto *anchor = getAnchor();
   auto path = getPath();
@@ -340,6 +350,10 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) {
       out << "key path dynamic member lookup";
       break;
 
+    case KeyPathType:
+      out << "key path type";
+      break;
+
     case KeyPathRoot:
       out << "key path root";
       break;

lib/Sema/ConstraintLocator.h

@@ -129,9 +129,11 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     SynthesizedArgument,
     /// The member looked up via keypath based dynamic lookup.
     KeyPathDynamicMember,
-    /// The root of a keypath
+    /// The type of the key path expression
+    KeyPathType,
+    /// The root of a key path
     KeyPathRoot,
-    /// The value of a keypath
+    /// The value of a key path
     KeyPathValue,
     /// The result type of a key path component. Not used for subscripts.
     KeyPathComponentResult,
@@ -165,6 +167,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ImplicitlyUnwrappedDisjunctionChoice:
     case DynamicLookupResult:
     case ContextualType:
+    case KeyPathType:
     case KeyPathRoot:
     case KeyPathValue:
     case KeyPathComponentResult:
@@ -234,6 +237,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ContextualType:
     case SynthesizedArgument:
     case KeyPathDynamicMember:
+    case KeyPathType:
     case KeyPathRoot:
     case KeyPathValue:
     case KeyPathComponentResult:
@@ -534,6 +538,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// e.g. `foo[0]` or `\Foo.[0]`
   bool isSubscriptMemberRef() const;
 
+  /// Determine whether give locator points to the type of the
+  /// key path expression.
+  bool isKeyPathType() const;
+
   /// Determine whether given locator points to the keypath root
   bool isKeyPathRoot() const;