[ConstraintSystem] Make sure that each index argument of keypath subscript is Hashable

xedin · xedin · commit 3288299e1e2b · 2019-04-10T13:50:34.000-07:00
Move checking for index Hashable conformance from CSApply (post-factum)
to the solver and use recorded conformance records complete subscript
components.
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -1561,15 +1561,6 @@ namespace {
           LocatorPathElt::getKeyPathDynamicMember(keyPathTy->getAnyNominal()));
       auto overload = solution.getOverloadChoice(componentLoc);
 
-      auto buildSubscriptComponent = [&](SourceLoc loc, Expr *indexExpr,
-                                         ArrayRef<Identifier> labels) {
-        // Save a reference to the component so we can do a post-pass to check
-        // the Hashable conformance of the indexes.
-        KeyPathSubscriptComponents.push_back({keyPath, 0});
-        return buildKeyPathSubscriptComponent(overload, loc, indexExpr, labels,
-                                              componentLoc);
-      };
-
       auto getKeyPathComponentIndex =
           [](ConstraintLocator *locator) -> unsigned {
         for (const auto &elt : locator->getPath()) {
@@ -1583,9 +1574,10 @@ namespace {
       // calls necessary to resolve a member reference.
       if (overload.choice.getKind() ==
           OverloadChoiceKind::KeyPathDynamicMemberLookup) {
-        keyPath->resolveComponents(
-            ctx, buildSubscriptComponent(dotLoc, /*indexExpr=*/nullptr,
-                                         ctx.Id_dynamicMember));
+        keyPath->resolveComponents(ctx,
+                                   buildKeyPathSubscriptComponent(
+                                       overload, dotLoc, /*indexExpr=*/nullptr,
+                                       ctx.Id_dynamicMember, componentLoc));
         return keyPath;
       }
 
@@ -1663,8 +1655,9 @@ namespace {
               /*implicit=*/true, SE->getAccessSemantics());
         }
 
-        component = buildSubscriptComponent(SE->getLoc(), SE->getIndex(),
-                                            SE->getArgumentLabels());
+        component = buildKeyPathSubscriptComponent(
+            overload, SE->getLoc(), SE->getIndex(), SE->getArgumentLabels(),
+            componentLoc);
       } else {
         return nullptr;
       }
@@ -4298,12 +4291,7 @@ namespace {
       E->setMethod(method);
       return E;
     }
-    
-  private:
-    // Key path components we need to
-    SmallVector<std::pair<KeyPathExpr *, unsigned>, 4>
-      KeyPathSubscriptComponents;
-  public:
+
     Expr *visitKeyPathExpr(KeyPathExpr *E) {
       if (E->isObjC()) {
         cs.setType(E, cs.getType(E->getObjCStringLiteralExpr()));
@@ -4445,10 +4433,6 @@ namespace {
               *foundDecl, origComponent.getLoc(), origComponent.getIndexExpr(),
               subscriptLabels, locator);
 
-          // Save a reference to the component so we can do a post-pass to check
-          // the Hashable conformance of the indexes.
-          KeyPathSubscriptComponents.push_back({E, resolvedComponents.size()});
-
           baseTy = component.getComponentType();
           resolvedComponents.push_back(component);
 
@@ -4624,10 +4608,13 @@ namespace {
       // through the subscript(dynamicMember:) member, restore the
       // openedType and origComponent to its full reference as if the user
       // wrote out the subscript manually.
-      if (overload.choice.getKind() ==
+      bool forDynamicLookup =
+          overload.choice.getKind() ==
               OverloadChoiceKind::DynamicMemberLookup ||
           overload.choice.getKind() ==
-              OverloadChoiceKind::KeyPathDynamicMemberLookup) {
+              OverloadChoiceKind::KeyPathDynamicMemberLookup;
+
+      if (forDynamicLookup) {
         overload.openedType =
             overload.openedFullType->castTo<AnyFunctionType>()->getResult();
 
@@ -4656,8 +4643,48 @@ namespace {
                               /*applyExpr*/ nullptr, labels,
                               /*hasTrailingClosure*/ false, locator);
 
-      return KeyPathExpr::Component::forSubscriptWithPrebuiltIndexExpr(
+      auto component = KeyPathExpr::Component::forSubscriptWithPrebuiltIndexExpr(
           ref, newIndexExpr, labels, resolvedTy, componentLoc, {});
+
+      // We need to be able to hash the captured index values in order for
+      // KeyPath itself to be hashable, so check that all of the subscript
+      // index components are hashable and collect their conformances here.
+      SmallVector<ProtocolConformanceRef, 4> conformances;
+
+      auto hashable =
+          cs.getASTContext().getProtocol(KnownProtocolKind::Hashable);
+
+      auto equatable =
+          cs.getASTContext().getProtocol(KnownProtocolKind::Equatable);
+
+      auto &TC = cs.getTypeChecker();
+      auto fnType = overload.openedType->castTo<FunctionType>();
+      for (const auto &param : fnType->getParams()) {
+        auto indexType = simplifyType(param.getPlainType());
+        // index conformance to the Hashable protocol has been verified
+        // by the solver, we just need to get it again with all of the
+        // generic parameters resolved.
+        auto hashableConformance =
+            TC.conformsToProtocol(indexType, hashable, cs.DC,
+                                  (ConformanceCheckFlags::Used |
+                                   ConformanceCheckFlags::InExpression));
+        assert(hashableConformance.hasValue());
+
+        // FIXME: Hashable implies Equatable, but we need to make sure the
+        // Equatable conformance is forced into existence during type
+        // checking so that it's available for SILGen.
+        auto eqConformance =
+            TC.conformsToProtocol(simplifyType(indexType), equatable, cs.DC,
+                                  (ConformanceCheckFlags::Used |
+                                   ConformanceCheckFlags::InExpression));
+        assert(eqConformance.hasValue());
+        (void)eqConformance;
+
+        conformances.push_back(*hashableConformance);
+      }
+
+      component.setSubscriptIndexHashableConformances(conformances);
+      return component;
     }
 
     Expr *visitKeyPathDotExpr(KeyPathDotExpr *E) {
@@ -4726,61 +4753,6 @@ namespace {
           .fixItInsertAfter(cast->getEndLoc(), ")");
       }
       
-      // Look at key path subscript components to verify that they're hashable.
-      for (auto componentRef : KeyPathSubscriptComponents) {
-        auto &component = componentRef.first
-                                  ->getMutableComponents()[componentRef.second];
-        // We need to be able to hash the captured index values in order for
-        // KeyPath itself to be hashable, so check that all of the subscript
-        // index components are hashable and collect their conformances here.
-        SmallVector<ProtocolConformanceRef, 2> hashables;
-        bool allIndexesHashable = true;
-        ArrayRef<TupleTypeElt> indexTypes;
-        TupleTypeElt singleIndexTypeBuf;
-        if (auto tup = cs.getType(component.getIndexExpr())
-                                               ->getAs<TupleType>()) {
-          indexTypes = tup->getElements();
-        } else {
-          singleIndexTypeBuf = cs.getType(component.getIndexExpr());
-          indexTypes = singleIndexTypeBuf;
-        }
-      
-        auto hashable =
-          cs.getASTContext().getProtocol(KnownProtocolKind::Hashable);
-        auto equatable =
-          cs.getASTContext().getProtocol(KnownProtocolKind::Equatable);
-        for (auto indexType : indexTypes) {
-          auto conformance =
-            cs.TC.conformsToProtocol(indexType.getType(), hashable,
-                                     cs.DC,
-                                     (ConformanceCheckFlags::Used|
-                                      ConformanceCheckFlags::InExpression));
-          if (!conformance) {
-            cs.TC.diagnose(component.getIndexExpr()->getLoc(),
-                           diag::expr_keypath_subscript_index_not_hashable,
-                           indexType.getType());
-            allIndexesHashable = false;
-            continue;
-          }
-          hashables.push_back(*conformance);
-          
-          // FIXME: Hashable implies Equatable, but we need to make sure the
-          // Equatable conformance is forced into existence during type checking
-          // so that it's available for SILGen.
-          auto eqConformance =
-            cs.TC.conformsToProtocol(indexType.getType(), equatable,
-                                     cs.DC,
-                                     (ConformanceCheckFlags::Used|
-                                      ConformanceCheckFlags::InExpression));
-          assert(eqConformance.hasValue());
-          (void)eqConformance;
-        }
-
-        if (allIndexesHashable) {
-          component.setSubscriptIndexHashableConformances(hashables);
-        }
-      }
-
       // Set the final types on the expression.
       cs.setExprTypes(result);
     }
diff --git a/lib/Sema/ConstraintLocator.h b/lib/Sema/ConstraintLocator.h
@@ -486,6 +486,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     bool isKeyPathDynamicMember() const {
       return getKind() == PathElementKind::KeyPathDynamicMember;
     }
+
+    bool isKeyPathComponent() const {
+      return getKind() == PathElementKind::KeyPathComponent;
+    }
   };
 
   /// Return the summary flags for an entire path.
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -1695,10 +1695,54 @@ isInvalidPartialApplication(ConstraintSystem &cs, const ValueDecl *member,
   return {true, level};
 }
 
+static bool isResultOfKeyPathDynamicMemberLookup(ConstraintLocator *locator) {
+  auto path = locator->getPath();
+  return !path.empty() && path.back().isKeyPathDynamicMember();
+}
+
+/// If given anchor + path represents a key path with subscript
+/// component at some index.
+static bool isKeyPathSubscriptComponent(ConstraintLocator *locator) {
+  auto *anchor = locator->getAnchor();
+  auto *KPE = dyn_cast_or_null<KeyPathExpr>(anchor);
+  if (!KPE)
+    return false;
+
+  using ComponentKind = KeyPathExpr::Component::Kind;
+  auto path = locator->getPath();
+  while (!path.empty()) {
+    auto &last = path.back();
+    if (!last.isKeyPathComponent()) {
+      path = path.drop_back();
+      continue;
+    }
+
+    auto index = last.getValue();
+    auto &component = KPE->getComponents()[index];
+    return component.getKind() == ComponentKind::Subscript ||
+           component.getKind() == ComponentKind::UnresolvedSubscript;
+  }
+  return false;
+}
+
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                        Type boundType,
                                        OverloadChoice choice,
                                        DeclContext *useDC) {
+  // Add a conformance constraint to make sure that given type conforms
+  // to Hashable protocol, which is important for key path subscript
+  // components.
+  auto verifyThatArgumentIsHashable = [&](unsigned index, Type argType,
+                                          ConstraintLocator *locator) {
+    if (auto *hashable = TC.getProtocol(choice.getDecl()->getLoc(),
+                                        KnownProtocolKind::Hashable)) {
+      addConstraint(ConstraintKind::ConformsTo, argType,
+                    hashable->getDeclaredType(),
+                    getConstraintLocator(
+                        locator, LocatorPathElt::getTupleElement(index)));
+    }
+  };
+
   // Determine the type to which we'll bind the overload set's type.
   Type refType;
   Type openedFullType;
@@ -1872,15 +1916,22 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
       assert(refFnType->getParams().size() == 1 &&
              "subscript always has one arg");
       auto argType = refFnType->getParams()[0].getPlainType();
-      
-      auto protoKind = KnownProtocolKind::ExpressibleByStringLiteral;
-      auto protocol = getTypeChecker().getProtocol(choice.getDecl()->getLoc(),
-                                                   protoKind);
-      if (!protocol)
+
+      auto &TC = getTypeChecker();
+
+      auto stringLiteral =
+          TC.getProtocol(choice.getDecl()->getLoc(),
+                         KnownProtocolKind::ExpressibleByStringLiteral);
+      if (!stringLiteral)
         break;
+
       addConstraint(ConstraintKind::LiteralConformsTo, argType,
-                    protocol->getDeclaredType(),
-                    locator);
+                    stringLiteral->getDeclaredType(), locator);
+
+      // If this is used inside of the keypath expression, we need to make
+      // sure that argument is Hashable.
+      if (isa<KeyPathExpr>(locator->getAnchor()))
+        verifyThatArgumentIsHashable(0, argType, locator);
     }
 
     if (kind == OverloadChoiceKind::KeyPathDynamicMemberLookup) {
@@ -1998,6 +2049,9 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
         // type into "leaf" directly.
         addConstraint(ConstraintKind::Equal, memberTy, leafTy, keyPathLoc);
       }
+
+      if (isa<KeyPathExpr>(locator->getAnchor()))
+        verifyThatArgumentIsHashable(0, keyPathTy, locator);
     }
     break;
   }
@@ -2067,6 +2121,21 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
       }
     }
 
+    if (auto *SD = dyn_cast<SubscriptDecl>(decl)) {
+      if (isResultOfKeyPathDynamicMemberLookup(locator) ||
+          isKeyPathSubscriptComponent(locator)) {
+        // Subscript type has a format of (Self[.Type) -> (Arg...) -> Result
+        auto declTy = openedFullType->castTo<FunctionType>();
+        auto subscriptTy = declTy->getResult()->castTo<FunctionType>();
+        // If we have subscript, each of the arguments has to conform to
+        // Hashable, because it would be used as a component inside key path.
+        for (auto index : indices(subscriptTy->getParams())) {
+          const auto &param = subscriptTy->getParams()[index];
+          verifyThatArgumentIsHashable(index, param.getPlainType(), locator);
+        }
+      }
+    }
+
     // Check whether applying this overload would result in invalid
     // partial function application e.g. partial application of
     // mutating method or initializer.
