[ConstraintSystem] Add KeyPathComponentTypes map

Avoids digging through the long list of type vars to find the ones related to key path components when dumping an expression.

Author: beccadax

lib/Sema/CSApply.cpp

@@ -1698,9 +1698,12 @@ namespace {
       }
 
       assert(componentExpr);
-      componentExpr->setType(simplifyType(cs.getType(anchor)));
+      Type ty = simplifyType(cs.getType(anchor));
+      componentExpr->setType(ty);
       cs.cacheType(componentExpr);
 
+      cs.setType(keyPath, 0, ty);
+
       keyPath->setParsedPath(componentExpr);
       keyPath->resolveComponents(ctx, {component});
       return keyPath;
@@ -4358,6 +4361,13 @@ namespace {
       Type baseTy = keyPathTy->getGenericArgs()[0];
       Type leafTy = keyPathTy->getGenericArgs()[1];
 
+      // Updates the constraint system with the type of the last resolved
+      // component. We do it this way because we sometimes insert new
+      // components.
+      auto updateCSWithResolvedComponent = [&]() {
+        cs.setType(E, resolvedComponents.size() - 1, baseTy);
+      };
+
       for (unsigned i : indices(E->getComponents())) {
         auto &origComponent = E->getMutableComponents()[i];
 
@@ -4434,6 +4444,7 @@ namespace {
           resolvedComponents.push_back(component);
 
           if (shouldForceUnwrapResult(foundDecl->choice, locator)) {
+            updateCSWithResolvedComponent();
             auto objectTy = getObjectType(baseTy);
             auto loc = origComponent.getLoc();
             component = KeyPathExpr::Component::forOptionalForce(objectTy, loc);
@@ -4461,6 +4472,7 @@ namespace {
           resolvedComponents.push_back(component);
 
           if (shouldForceUnwrapResult(foundDecl->choice, locator)) {
+            updateCSWithResolvedComponent();
             auto objectTy = getObjectType(baseTy);
             auto loc = origComponent.getLoc();
             component = KeyPathExpr::Component::forOptionalForce(objectTy, loc);
@@ -4512,6 +4524,9 @@ namespace {
         case KeyPathExpr::Component::Kind::TupleElement:
           llvm_unreachable("already resolved");
         }
+
+        // By now, "baseTy" is the result type of this component.
+        updateCSWithResolvedComponent();
       }
 
       // Wrap a non-optional result if there was chaining involved.
@@ -4524,6 +4539,7 @@ namespace {
         auto component = KeyPathExpr::Component::forOptionalWrap(leafTy);
         resolvedComponents.push_back(component);
         baseTy = leafTy;
+        updateCSWithResolvedComponent();
       }
       E->resolveComponents(cs.getASTContext(), resolvedComponents);
 

lib/Sema/CSGen.cpp

@@ -3054,6 +3054,10 @@ namespace {
         case KeyPathExpr::Component::Kind::Identity:
           continue;
         }
+
+        // By now, `base` is the result type of this component. Set it in the
+        // constraint system so we can find it later.
+        CS.setType(E, i, base);
       }
 
       // If there was an optional chaining component, the end result must be

lib/Sema/ConstraintSystem.h

@@ -1065,6 +1065,8 @@ class ConstraintSystem {
   llvm::DenseMap<const Expr *, TypeBase *> ExprTypes;
   llvm::DenseMap<const TypeLoc *, TypeBase *> TypeLocTypes;
   llvm::DenseMap<const ParamDecl *, TypeBase *> ParamTypes;
+  llvm::DenseMap<std::pair<const KeyPathExpr *, unsigned>, TypeBase *>
+      KeyPathComponentTypes;
 
   /// Maps closure parameters to type variables.
   llvm::DenseMap<const ParamDecl *, TypeVariableType *>
@@ -1446,6 +1448,11 @@ class ConstraintSystem {
       if (expr->getType())
         CS.cacheType(expr);
 
+      if (auto kp = dyn_cast<KeyPathExpr>(expr))
+        for (auto i : indices(kp->getComponents()))
+          if (kp->getComponents()[i].getComponentType())
+            CS.cacheType(kp, i);
+
       return expr;
     }
 
@@ -1477,6 +1484,15 @@ class ConstraintSystem {
              "Should not write type variable into expression!");
       expr->setType(CS.getType(expr));
 
+      if (auto kp = dyn_cast<KeyPathExpr>(expr)) {
+        for (auto i : indices(kp->getComponents())) {
+          Type componentType;
+          if (CS.hasType(kp, i))
+            componentType = CS.getType(kp, i);
+          kp->getMutableComponents()[i].setComponentType(componentType);
+        }
+      }
+
       return expr;
     }
 
@@ -1757,6 +1773,12 @@ class ConstraintSystem {
     ParamTypes[P] = T.getPointer();
   }
 
+  void setType(KeyPathExpr *KP, unsigned I, Type T) {
+    assert(KP && "Expected non-null key path parameter!");
+    assert(T && "Expected non-null type!");
+    KeyPathComponentTypes[std::make_pair(KP, I)] = T.getPointer();
+  }
+
   /// Check to see if we have a type for an expression.
   bool hasType(const Expr *E) const {
     assert(E != nullptr && "Expected non-null expression!");
@@ -1772,6 +1794,12 @@ class ConstraintSystem {
     return ParamTypes.find(P) != ParamTypes.end();
   }
 
+  bool hasType(const KeyPathExpr *KP, unsigned I) const {
+    assert(KP && "Expected non-null key path parameter!");
+    return KeyPathComponentTypes.find(std::make_pair(KP, I))
+              != KeyPathComponentTypes.end();
+  }
+
   /// Get the type for an expression.
   Type getType(const Expr *E) const {
     assert(hasType(E) && "Expected type to have been set!");
@@ -1800,6 +1828,11 @@ class ConstraintSystem {
     return wantInterfaceType ? D->getInterfaceType() : D->getType();
   }
 
+  Type getType(const KeyPathExpr *KP, unsigned I) const {
+    assert(hasType(KP, I) && "Expected type to have been set!");
+    return KeyPathComponentTypes.find(std::make_pair(KP, I))->second;
+  }
+
   /// Cache the type of the expression argument and return that same
   /// argument.
   template <typename T>
@@ -1809,6 +1842,15 @@ class ConstraintSystem {
     return E;
   }
 
+  /// Cache the type of the expression argument and return that same
+  /// argument.
+  KeyPathExpr *cacheType(KeyPathExpr *E, unsigned I) {
+    auto componentTy = E->getComponents()[I].getComponentType();
+    assert(componentTy && "Expected a type!");
+    setType(E, I, componentTy);
+    return E;
+  }
+
   void setContextualType(Expr *E, TypeLoc T, ContextualTypePurpose purpose) {
     assert(E != nullptr && "Expected non-null expression!");
     contextualTypeNode = E;

lib/Sema/TypeCheckConstraints.cpp

@@ -3408,27 +3408,9 @@ void ConstraintSystem::print(raw_ostream &out, Expr *E) {
     return Type();
   };
   auto getTypeOfKeyPathComponent =
-      [&](const KeyPathExpr *E, unsigned index) -> Type {
-    // CSGen attaches the type var for each key path component to the
-    // "key path component #n -> function result" locator. This is the natural
-    // place subscript result types will end up, and it's otherwise unused by
-    // other components.
-    auto componentLocator = getConstraintLocator(E,
-        ConstraintLocator::PathElement::getKeyPathComponent(index));
-    auto resultLocator = getConstraintLocator(componentLocator,
-        ConstraintLocator::FunctionResult);
-
-    // Find the first type variable with this locator and return it.
-    for (auto &typeVar : TypeVariables) {
-      if ((*typeVar)->getLocator() != resultLocator)
-        continue;
-
-      if (auto fixedTy = (*typeVar)->getFixedType(getSavedBindings()))
-        return fixedTy;
-
-      return typeVar;
-    }
-
+      [&](const KeyPathExpr *KP, unsigned I) -> Type {
+    if (hasType(KP, I))
+      return getType(KP, I);
     return Type();
   };