Merge pull request #67663 from amritpan/kp-root-value-tv-inference

[ConstraintSystem] Allow inference to bind AnyKeyPath as a KeyPath

Author: xedin

include/swift/Sema/ConstraintLocator.h

@@ -1062,22 +1062,6 @@ 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;
-  }
-};
-
 class LocatorPathElt::ConstructorMemberType final
     : public StoredIntegerElement<1> {
 public:

include/swift/Sema/ConstraintLocatorPathElts.def

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

include/swift/Sema/ConstraintSystem.h

@@ -3031,6 +3031,20 @@ class ConstraintSystem {
       return std::get<1>(result->second);
     return nullptr;
   }
+  
+  TypeVariableType *getKeyPathRootType(const KeyPathExpr *keyPath) const {
+    auto result = getKeyPathRootTypeIfAvailable(keyPath);
+    assert(result);
+    return result;
+  }
+
+  TypeVariableType *
+  getKeyPathRootTypeIfAvailable(const KeyPathExpr *keyPath) const {
+    auto result = KeyPaths.find(keyPath);
+    if (result != KeyPaths.end())
+      return std::get<0>(result->second);
+    return nullptr;
+  }
 
   TypeBase* getFavoredType(Expr *E) {
     assert(E != nullptr);

lib/Sema/CSBindings.cpp

@@ -1250,27 +1250,34 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
     return llvm::None;
 
   if (TypeVar->getImpl().isKeyPathType()) {
-    auto *BGT = type->lookThroughAllOptionalTypes()->getAs<BoundGenericType>();
-    if (!BGT || !isKnownKeyPathType(BGT))
+    auto objectTy = type->lookThroughAllOptionalTypes();
+    if (!isKnownKeyPathType(objectTy))
       return llvm::None;
-
-    // `PartialKeyPath<T>` represents a type-erased version of `KeyPath<T, V>`.
+    
+    auto &ctx = CS.getASTContext();
+    auto *keyPathTypeLoc = TypeVar->getImpl().getLocator();
+    auto *keyPath = castToExpr<KeyPathExpr>(keyPathTypeLoc->getAnchor());
+    // `AnyKeyPath` and `PartialKeyPath` represent type-erased versions of
+    // `KeyPath<T, V>`.
     //
-    // In situations where partial key path cannot be used directly i.e.
-    // passing an argument to a parameter represented by a partial key path,
-    // let's attempt a `KeyPath` binding which would then be converted to a
-    // partial key path since there is a subtype relationship between them.
-    if (BGT->isPartialKeyPath() && kind == AllowedBindingKind::Subtypes) {
-      auto &ctx = CS.getASTContext();
-      auto *keyPathLoc = TypeVar->getImpl().getLocator();
-
-      auto rootTy = BGT->getGenericArgs()[0];
+    // In situations where `AnyKeyPath` or `PartialKeyPath` cannot be used
+    // directly i.e. passing an argument to a parameter represented by a
+    // `AnyKeyPath` or `PartialKeyPath`, let's attempt a `KeyPath` binding which
+    // would then be converted to a `AnyKeyPath` or `PartialKeyPath` since there
+    // is a subtype relationship between them.
+    if (objectTy->isAnyKeyPath()) {
+      auto root = CS.getKeyPathRootType(keyPath);
+      auto value = CS.getKeyPathValueType(keyPath);
+      
+      type = BoundGenericType::get(ctx.getKeyPathDecl(), Type(),
+                                   {root, value});
+    } else if (objectTy->isPartialKeyPath() &&
+               kind == AllowedBindingKind::Subtypes) {
+      auto rootTy = objectTy->castTo<BoundGenericType>()->getGenericArgs()[0];
       // Since partial key path is an erased version of `KeyPath`, the value
       // type would never be used, which means that binding can use
       // type variable generated for a result of key path expression.
-      auto valueTy =
-          keyPathLoc->castLastElementTo<LocatorPathElt::KeyPathType>()
-              .getValueType();
+      auto valueTy = CS.getKeyPathValueType(keyPath);
 
       type = BoundGenericType::get(ctx.getKeyPathDecl(), Type(),
                                    {rootTy, valueTy});

lib/Sema/CSGen.cpp

@@ -3822,7 +3822,7 @@ namespace {
       // The type of key path depends on the overloads chosen for the key
       // path components.
       auto typeLoc =
-          CS.getConstraintLocator(locator, LocatorPathElt::KeyPathType(value));
+          CS.getConstraintLocator(locator, LocatorPathElt::KeyPathType());
 
       Type kpTy = CS.createTypeVariable(typeLoc, TVO_CanBindToNoEscape |
                                                      TVO_CanBindToHole);