Add root type constraint between KeyPath expressions and applications

lib/Sema/CSSimplify.cpp

@@ -4955,11 +4955,50 @@ ConstraintSystem::addConstraintImpl(ConstraintKind kind, Type first,
   llvm_unreachable("Unhandled ConstraintKind in switch.");
 }
 
+void
+ConstraintSystem::addKeyPathApplicationRootConstraint(Type root, ConstraintLocatorBuilder locator) {
+  // If this is a subscript with a KeyPath expression, add a constraint that
+  // connects the subscript's root type to the root type of the KeyPath.
+  SmallVector<LocatorPathElt, 4> path;
+  Expr *anchor = locator.getLocatorParts(path);
+
+  auto subscript = dyn_cast_or_null<SubscriptExpr>(anchor);
+  if (!subscript)
+    return;
+
+  assert(path.size() == 1 && path[0].getKind() == ConstraintLocator::SubscriptMember);
+  auto indexTuple = dyn_cast<TupleExpr>(subscript->getIndex());
+  if (!indexTuple || indexTuple->getNumElements() != 1)
+    return;
+
+  auto keyPathExpr = dyn_cast<KeyPathExpr>(indexTuple->getElement(0));
+  if (!keyPathExpr)
+    return;
+
+  auto typeVar = getType(keyPathExpr)->getAs<TypeVariableType>();
+  if (!typeVar)
+    return;
+
+  SmallVector<Constraint *, 4> constraints;
+  CG.gatherConstraints(typeVar, constraints,
+                       ConstraintGraph::GatheringKind::EquivalenceClass);
+
+  for (auto constraint : constraints) {
+    if (constraint->getKind() == ConstraintKind::KeyPath &&
+        constraint->getLocator()->getAnchor() == keyPathExpr) {
+      auto keyPathRootTy = constraint->getSecondType();
+      addConstraint(ConstraintKind::Subtype, root->getWithoutSpecifierType(), keyPathRootTy, locator);
+    }
+  }
+}
+
 void
 ConstraintSystem::addKeyPathApplicationConstraint(Type keypath,
                                               Type root, Type value,
                                               ConstraintLocatorBuilder locator,
                                               bool isFavored) {
+  addKeyPathApplicationRootConstraint(root, locator);
+
   switch (simplifyKeyPathApplicationConstraint(keypath, root, value,
                                                TMF_GenerateConstraints,
                                                locator)) {

lib/Sema/ConstraintSystem.h

@@ -1553,6 +1553,10 @@ class ConstraintSystem {
 
   /// Add a new type variable that was already created.
   void addTypeVariable(TypeVariableType *typeVar);
+
+  /// \brief Add a constraint from the subscript base to the root of the key
+  /// path literal to the constraint system.
+  void addKeyPathApplicationRootConstraint(Type root, ConstraintLocatorBuilder locator);
 
 public:
   /// \brief Lookup for a member with the given name in the given base type.

test/expr/unary/keypath/keypath.swift

@@ -499,6 +499,41 @@ func sr6744() {
     _ = get(for: \.value)
 }
 
+func sr7380() {
+  _ = ""[keyPath: \.count]
+  _ = ""[keyPath: \String.count]
+
+  let arr1 = [1]
+  _ = arr1[keyPath: \.[0]]
+  _ = arr1[keyPath: \[Int].[0]]
+
+  let dic1 = [1:"s"]
+  _ = dic1[keyPath: \.[1]]
+  _ = dic1[keyPath: \[Int: String].[1]]
+
+  var arr2 = [1]
+  arr2[keyPath: \.[0]] = 2
+  arr2[keyPath: \[Int].[0]] = 2
+
+  var dic2 = [1:"s"]
+  dic2[keyPath: \.[1]] = ""
+  dic2[keyPath: \[Int: String].[1]] = ""
+
+  _ = [""][keyPath: \.[0]]
+  _ = [""][keyPath: \[String].[0]]
+
+  _ = ["": ""][keyPath: \.["foo"]]
+  _ = ["": ""][keyPath: \[String: String].["foo"]]
+
+  class A {
+    var a: String = ""
+  }
+  _ = A()[keyPath: \.a]
+  _ = A()[keyPath: \A.a]
+  A()[keyPath: \.a] = ""
+  A()[keyPath: \A.a] = ""
+}
+
 struct VisibilityTesting {
   private(set) var x: Int
   fileprivate(set) var y: Int