Merge pull request #39612 from Jumhyn/keypath-function-conversion

Jumhyn · web-flow · commit 0735629c8da0 · 2024-02-03T10:26:13.000-08:00
[ConstraintSystem] Allow function-function conversions for keypath literals
diff --git a/include/swift/AST/DiagnosticsSema.def b/include/swift/AST/DiagnosticsSema.def
@@ -697,6 +697,12 @@ ERROR(expr_smart_keypath_application_type_mismatch,none,
       "key path of type %0 cannot be applied to a base of type %1",
       (Type, Type))
 ERROR(expr_keypath_root_type_mismatch, none,
+      "key path root type %0 cannot be converted to contextual type %1",
+      (Type, Type))
+ERROR(expr_keypath_type_mismatch, none,
+      "key path of type %0 cannot be converted to contextual type %1",
+      (Type, Type))
+ERROR(expr_keypath_application_root_type_mismatch, none,
       "key path with root type %0 cannot be applied to a base of type %1",
       (Type, Type))
 ERROR(expr_swift_keypath_anyobject_root,none,
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -1737,6 +1737,11 @@ class Solution {
   /// Retrieve the type of the \p ComponentIndex-th component in \p KP.
   Type getType(const KeyPathExpr *KP, unsigned ComponentIndex) const;
 
+  TypeVariableType *getKeyPathRootType(const KeyPathExpr *keyPath) const;
+
+  TypeVariableType *
+  getKeyPathRootTypeIfAvailable(const KeyPathExpr *keyPath) const;
+
   /// Retrieve the type of the given node as recorded in this solution
   /// and resolve all of the type variables in contains to form a fully
   /// "resolved" concrete type.
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -4971,20 +4971,18 @@ namespace {
       // Resolve each of the components.
       bool didOptionalChain = false;
       bool isFunctionType = false;
-      Type baseTy, leafTy;
+      auto baseTy = cs.simplifyType(solution.getKeyPathRootType(E));
+      Type leafTy;
       Type exprType = cs.getType(E);
       if (auto fnTy = exprType->getAs<FunctionType>()) {
-        baseTy = fnTy->getParams()[0].getParameterType();
         leafTy = fnTy->getResult();
         isFunctionType = true;
       } else if (auto *existential = exprType->getAs<ExistentialType>()) {
         auto layout = existential->getExistentialLayout();
         auto keyPathTy = layout.explicitSuperclass->castTo<BoundGenericType>();
-        baseTy = keyPathTy->getGenericArgs()[0];
         leafTy = keyPathTy->getGenericArgs()[1];
       } else {
         auto keyPathTy = exprType->castTo<BoundGenericType>();
-        baseTy = keyPathTy->getGenericArgs()[0];
         leafTy = keyPathTy->getGenericArgs()[1];
       }
 
@@ -5103,13 +5101,11 @@ namespace {
         assert(!resolvedComponents.empty());
         componentTy = resolvedComponents.back().getComponentType();
       }
-      
+
       // Wrap a non-optional result if there was chaining involved.
       if (didOptionalChain && componentTy &&
           !componentTy->hasUnresolvedType() &&
           !componentTy->getWithoutSpecifierType()->isEqual(leafTy)) {
-        assert(leafTy->getOptionalObjectType()->isEqual(
-            componentTy->getWithoutSpecifierType()));
         auto component = KeyPathExpr::Component::forOptionalWrap(leafTy);
         resolvedComponents.push_back(component);
         componentTy = leafTy;
@@ -5122,11 +5118,6 @@ namespace {
       // See whether there's an equivalent ObjC key path string we can produce
       // for interop purposes.
       checkAndSetObjCKeyPathString(E);
-      
-      // The final component type ought to line up with the leaf type of the
-      // key path.
-      assert(!componentTy || componentTy->hasUnresolvedType()
-             || componentTy->getWithoutSpecifierType()->isEqual(leafTy));
 
       if (!isFunctionType)
         return E;
@@ -9848,6 +9839,21 @@ Type Solution::getType(const KeyPathExpr *KP, unsigned I) const {
   return keyPathComponentTypes.find(std::make_pair(KP, I))->second;
 }
 
+TypeVariableType *
+Solution::getKeyPathRootType(const KeyPathExpr *keyPath) const {
+  auto result = getKeyPathRootTypeIfAvailable(keyPath);
+  assert(result);
+  return result;
+}
+
+TypeVariableType *
+Solution::getKeyPathRootTypeIfAvailable(const KeyPathExpr *keyPath) const {
+  auto result = KeyPaths.find(keyPath);
+  if (result != KeyPaths.end())
+    return std::get<0>(result->second);
+  return nullptr;
+}
+
 Type Solution::getResolvedType(ASTNode node) const {
   return simplifyType(getType(node));
 }
diff --git a/lib/Sema/CSDiagnostics.cpp b/lib/Sema/CSDiagnostics.cpp
@@ -2497,9 +2497,13 @@ bool ContextualFailure::diagnoseAsError() {
 
   if (path.empty()) {
     if (auto *KPE = getAsExpr<KeyPathExpr>(anchor)) {
-      emitDiagnosticAt(KPE->getLoc(),
-                       diag::expr_keypath_type_covert_to_contextual_type,
-                       getFromType(), getToType());
+      Diag<Type, Type> diag;
+      if (auto ctxDiag = getDiagnosticFor(CTP, getToType())) {
+        diag = *ctxDiag;
+      } else {
+        diag = diag::expr_keypath_type_mismatch;
+      }
+      emitDiagnosticAt(KPE->getLoc(), diag, getFromType(), getToType());
       return true;
     }
 
@@ -2750,9 +2754,14 @@ bool ContextualFailure::diagnoseAsError() {
     break;
   }
 
+  case ConstraintLocator::FunctionResult:
   case ConstraintLocator::KeyPathValue: {
-    diagnostic = diag::expr_keypath_value_covert_to_contextual_type;
-    break;
+    if (auto *KPE = getAsExpr<KeyPathExpr>(anchor)) {
+      diagnostic = diag::expr_keypath_value_covert_to_contextual_type;
+      break;
+    } else {
+      return false;
+    }
   }
 
   default:
@@ -8272,13 +8281,24 @@ bool CoercionAsForceCastFailure::diagnoseAsError() {
 
 bool KeyPathRootTypeMismatchFailure::diagnoseAsError() {
   auto locator = getLocator();
+  auto anchor = locator->getAnchor();
   assert(locator->isKeyPathRoot() && "Expected a key path root");
-  
-  auto baseType = getFromType();
-  auto rootType = getToType();
 
-  emitDiagnostic(diag::expr_keypath_root_type_mismatch,
-                 rootType, baseType);
+
+
+  if (isExpr<KeyPathApplicationExpr>(anchor) || isExpr<SubscriptExpr>(anchor)) {
+    auto baseType = getFromType();
+    auto rootType = getToType();
+
+    emitDiagnostic(diag::expr_keypath_application_root_type_mismatch,
+                   rootType, baseType);
+  } else {
+    auto rootType = getFromType();
+    auto expectedType = getToType();
+
+    emitDiagnostic(diag::expr_keypath_root_type_mismatch, rootType,
+                   expectedType);
+  }
   return true;
 }
 
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -5222,6 +5222,13 @@ bool ConstraintSystem::repairFailures(
                         });
   };
 
+  auto hasAnyRestriction = [&]() {
+    return llvm::any_of(conversionsOrFixes,
+                        [](const RestrictionOrFix &correction) {
+                          return bool(correction.getRestriction());
+                        });
+  };
+
   // Check whether this is a tuple with a single unlabeled element
   // i.e. `(_: Int)` and return type of that element if so. Note that
   // if the element is pack expansion type the tuple is significant.
@@ -5247,6 +5254,40 @@ bool ConstraintSystem::repairFailures(
     return true;
   }
 
+  auto maybeRepairKeyPathResultFailure = [&](KeyPathExpr *kpExpr) {
+    if (lhs->isPlaceholder() || rhs->isPlaceholder())
+      return true;
+    if (lhs->isTypeVariableOrMember() || rhs->isTypeVariableOrMember())
+      return false;
+
+    if (hasConversionOrRestriction(ConversionRestrictionKind::DeepEquality) ||
+        hasConversionOrRestriction(ConversionRestrictionKind::ValueToOptional))
+      return false;
+
+    auto i = kpExpr->getComponents().size() - 1;
+    auto lastCompLoc =
+        getConstraintLocator(kpExpr, LocatorPathElt::KeyPathComponent(i));
+    if (hasFixFor(lastCompLoc, FixKind::AllowTypeOrInstanceMember))
+      return true;
+
+    auto *keyPathLoc = getConstraintLocator(anchor);
+
+    if (hasFixFor(keyPathLoc))
+      return true;
+
+    if (auto contextualInfo = getContextualTypeInfo(anchor)) {
+      if (hasFixFor(getConstraintLocator(
+              keyPathLoc,
+              LocatorPathElt::ContextualType(contextualInfo->purpose))))
+        return true;
+    }
+
+    conversionsOrFixes.push_back(IgnoreContextualType::create(
+        *this, lhs, rhs,
+        getConstraintLocator(keyPathLoc, ConstraintLocator::KeyPathValue)));
+    return true;
+  };
+
   if (path.empty()) {
     if (!anchor)
       return false;
@@ -5266,9 +5307,9 @@ bool ConstraintSystem::repairFailures(
     // instance fix recorded.
     if (auto *kpExpr = getAsExpr<KeyPathExpr>(anchor)) {
       if (isKnownKeyPathType(lhs) && isKnownKeyPathType(rhs)) {
-        // If we have keypath capabilities for both sides and one of the bases
-        // is unresolved, it is too early to record fix.
-        if (hasConversionOrRestriction(ConversionRestrictionKind::DeepEquality))
+        // If we have a conversion happening here, we should let fix happen in
+        // simplifyRestrictedConstraint.
+        if (hasAnyRestriction())
           return false;
       }
 
@@ -5668,10 +5709,7 @@ bool ConstraintSystem::repairFailures(
 
     // If there are any restrictions here we need to wait and let
     // `simplifyRestrictedConstraintImpl` handle them.
-    if (llvm::any_of(conversionsOrFixes,
-                     [](const RestrictionOrFix &correction) {
-                       return bool(correction.getRestriction());
-                     }))
+    if (hasAnyRestriction())
       break;
 
     if (auto *fix = fixPropertyWrapperFailure(
@@ -6090,10 +6128,7 @@ bool ConstraintSystem::repairFailures(
 
     // If there are any restrictions here we need to wait and let
     // `simplifyRestrictedConstraintImpl` handle them.
-    if (llvm::any_of(conversionsOrFixes,
-                     [](const RestrictionOrFix &correction) {
-                       return bool(correction.getRestriction());
-                     }))
+    if (hasAnyRestriction())
       break;
 
     // `lhs` - is an result type and `rhs` is a contextual type.
@@ -6112,6 +6147,10 @@ bool ConstraintSystem::repairFailures(
       return true;
     }
 
+    if (auto *kpExpr = getAsExpr<KeyPathExpr>(anchor)) {
+      return maybeRepairKeyPathResultFailure(kpExpr);
+    }
+
     auto *loc = getConstraintLocator(anchor, {path.begin(), path.end() - 1});
     // If this is a mismatch between contextual type and (trailing)
     // closure with explicitly specified result type let's record it
@@ -6683,37 +6722,9 @@ bool ConstraintSystem::repairFailures(
     return true;
   }
   case ConstraintLocator::KeyPathValue: {
-    if (lhs->isPlaceholder() || rhs->isPlaceholder())
-      return true;
-    if (lhs->isTypeVariableOrMember() || rhs->isTypeVariableOrMember())
-      break;
-
-    if (hasConversionOrRestriction(ConversionRestrictionKind::DeepEquality) ||
-        hasConversionOrRestriction(ConversionRestrictionKind::ValueToOptional))
-      return false;
-
-    auto kpExpr = castToExpr<KeyPathExpr>(anchor);
-    auto i = kpExpr->getComponents().size() - 1;
-    auto lastCompLoc =
-        getConstraintLocator(kpExpr, LocatorPathElt::KeyPathComponent(i));
-    if (hasFixFor(lastCompLoc, FixKind::AllowTypeOrInstanceMember))
+    if (maybeRepairKeyPathResultFailure(getAsExpr<KeyPathExpr>(anchor)))
       return true;
 
-    auto *keyPathLoc = getConstraintLocator(anchor);
-
-    if (hasFixFor(keyPathLoc))
-      return true;
-
-    if (auto contextualInfo = getContextualTypeInfo(anchor)) {
-      if (hasFixFor(getConstraintLocator(
-              keyPathLoc,
-              LocatorPathElt::ContextualType(contextualInfo->purpose))))
-        return true;
-    }
-
-    conversionsOrFixes.push_back(IgnoreContextualType::create(
-        *this, lhs, rhs,
-        getConstraintLocator(keyPathLoc, ConstraintLocator::KeyPathValue)));
     break;
   }
   default:
@@ -12257,12 +12268,26 @@ ConstraintSystem::simplifyKeyPathConstraint(
 
     if (auto fnTy = contextualTy->getAs<FunctionType>()) {
       assert(fnTy->getParams().size() == 1);
-      // Match up the root and value types to the function's param and return
-      // types. Note that we're using the type of the parameter as referenced
-      // from inside the function body as we'll be transforming the code into:
-      // { root in root[keyPath: kp] }.
-      contextualRootTy = fnTy->getParams()[0].getParameterType();
-      contextualValueTy = fnTy->getResult();
+      // Key paths may be converted to a function of compatible type. We will
+      // later form from this key path an implicit closure of the form
+      // `{ root in root[keyPath: kp] }` so any conversions that are valid with
+      // a source type of `(Root) -> Value` should be valid here too.
+      auto rootParam = AnyFunctionType::Param(rootTy);
+      auto kpFnTy = FunctionType::get(rootParam, valueTy, fnTy->getExtInfo());
+
+      // Note: because the keypath is applied to `root` as a parameter internal
+      // to the closure, we use the function parameter's "parameter type" rather
+      // than the raw type. This enables things like:
+      // ```
+      // let countKeyPath: (String...) -> Int = \.count
+      // ```
+      auto paramTy = fnTy->getParams()[0].getParameterType();
+      auto paramParam = AnyFunctionType::Param(paramTy);
+      auto paramFnTy = FunctionType::get(paramParam, fnTy->getResult(),
+                                         fnTy->getExtInfo());
+
+      return matchTypes(kpFnTy, paramFnTy, ConstraintKind::Conversion, subflags,
+                        locator).isSuccess();
     }
 
     assert(contextualRootTy && contextualValueTy);
diff --git a/test/Constraints/keypath.swift b/test/Constraints/keypath.swift
@@ -81,8 +81,8 @@ func testVariadicKeypathAsFunc() {
 
   // These are not okay, the KeyPath should have a base that matches the
   // internal parameter type of the function, i.e (S...).
-  let _: (S...) -> Int = \S.i // expected-error {{key path with root type 'S...' cannot be applied to a base of type 'S'}}
-  takesVariadicFnWithGenericRet(\S.i) // expected-error {{key path with root type 'S...' cannot be applied to a base of type 'S'}}
+  let _: (S...) -> Int = \S.i // expected-error {{key path root type 'S' cannot be converted to contextual type 'S...'}}
+  takesVariadicFnWithGenericRet(\S.i) // expected-error {{key path root type 'S' cannot be converted to contextual type 'S...'}}
 }
 
 // rdar://problem/54322807
@@ -231,7 +231,7 @@ func issue_65965() {
   let refKP: ReferenceWritableKeyPath<S, String>
   refKP = \.s
   // expected-error@-1 {{cannot convert key path type 'WritableKeyPath<S, String>' to contextual type 'ReferenceWritableKeyPath<S, String>'}}
-	
+
   let writeKP: WritableKeyPath<S, String>
   writeKP = \.v
   // expected-error@-1 {{cannot convert key path type 'KeyPath<S, String>' to contextual type 'WritableKeyPath<S, String>'}}
diff --git a/test/expr/unary/keypath/keypath.swift b/test/expr/unary/keypath/keypath.swift
diff --git a/unittests/Sema/CMakeLists.txt b/unittests/Sema/CMakeLists.txt
diff --git a/unittests/Sema/KeypathFunctionConversionTests.cpp b/unittests/Sema/KeypathFunctionConversionTests.cpp
