Merge pull request #70007 from xedin/keypath-application-improvements

[ConstraintSystem] Modernization of key path application handling

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -704,6 +704,10 @@ ERROR(expr_swift_keypath_anyobject_root,none,
       "the root type of a Swift key path cannot be 'AnyObject'", ())
 ERROR(expr_keypath_multiparam_func_conversion, none,
       "cannot convert key path into a multi-argument function type %0", (Type))
+ERROR(cannot_convert_type_to_keypath_subscript_index,none,
+      "cannot use value of type %0 as a key path subscript index; argument must be"
+      " a key path",
+      (Type))
 WARNING(expr_deprecated_writable_keypath,Deprecation,
         "forming a writable keypath to property %0 that is read-only in this context "
         "is deprecated and will be removed in a future release",

include/swift/Sema/CSFix.h

@@ -463,6 +463,10 @@ enum class FixKind : uint8_t {
   /// Ignore situations when provided number of generic arguments didn't match
   /// expected number of parameters.
   IgnoreGenericSpecializationArityMismatch,
+
+  /// Ignore situations when key path subscript index gets passed an invalid
+  /// type as an argument (something that is not a key path).
+  IgnoreKeyPathSubscriptIndexMismatch,
 };
 
 class ConstraintFix {
@@ -3708,6 +3712,30 @@ class IgnoreGenericSpecializationArityMismatch final : public ConstraintFix {
   }
 };
 
+class IgnoreKeyPathSubscriptIndexMismatch final : public ConstraintFix {
+  Type ArgType;
+
+  IgnoreKeyPathSubscriptIndexMismatch(ConstraintSystem &cs, Type argTy,
+                                      ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::IgnoreKeyPathSubscriptIndexMismatch,
+                      locator),
+        ArgType(argTy) {}
+
+public:
+  std::string getName() const override {
+    return "ignore invalid key path subscript index";
+  }
+
+  bool diagnose(const Solution &solution, bool asNote = false) const override;
+
+  static IgnoreKeyPathSubscriptIndexMismatch *
+  create(ConstraintSystem &cs, Type argType, ConstraintLocator *locator);
+
+  static bool classof(const ConstraintFix *fix) {
+    return fix->getKind() == FixKind::IgnoreKeyPathSubscriptIndexMismatch;
+  }
+};
+
 } // end namespace constraints
 } // end namespace swift
 

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -131,6 +131,9 @@ SIMPLE_LOCATOR_PATH_ELT(KeyPathType)
 /// The value of a key path.
 SIMPLE_LOCATOR_PATH_ELT(KeyPathValue)
 
+/// The index parameter of a special `subscript(keyPath: KeyPath<...>)`.
+SIMPLE_LOCATOR_PATH_ELT(KeyPathSubscriptIndex)
+
 /// An implicit @lvalue-to-inout conversion; only valid for operator
 /// arguments.
 SIMPLE_LOCATOR_PATH_ELT(LValueConversion)

include/swift/Sema/ConstraintSystem.h

@@ -499,6 +499,10 @@ class TypeVariableType::Implementation {
   /// expression.
   bool isKeyPathValue() const;
 
+  /// Determine whether this type variable represents an index parameter of
+  /// a special `subscript(keyPath:)` subscript.
+  bool isKeyPathSubscriptIndex() const;
+
   /// Determine whether this type variable represents a subscript result type.
   bool isSubscriptResultType() const;
 

lib/Sema/CSBindings.cpp

@@ -1427,6 +1427,27 @@ PotentialBindings::inferFromRelational(Constraint *constraint) {
       return llvm::None;
   }
 
+  if (TypeVar->getImpl().isKeyPathSubscriptIndex()) {
+    // Key path subscript index can only be a r-value non-optional
+    // type that is a subtype of a known KeyPath type.
+    type = type->getRValueType()->lookThroughAllOptionalTypes();
+
+    // If argument to a key path subscript is an existential,
+    // we can erase it to superclass (if any) here and solver
+    // will perform the opening if supertype turns out to be
+    // a valid key path type of its subtype.
+    if (kind == AllowedBindingKind::Supertypes) {
+      if (type->isExistentialType()) {
+        auto layout = type->getExistentialLayout();
+        if (auto superclass = layout.explicitSuperclass) {
+          type = superclass;
+        } else if (!CS.shouldAttemptFixes()) {
+          return llvm::None;
+        }
+      }
+    }
+  }
+
   if (auto *locator = TypeVar->getImpl().getLocator()) {
     // Don't allow a protocol type to get propagated from the base to the result
     // type of a chain, Result should always be a concrete type which conforms

lib/Sema/CSDiagnostics.cpp

@@ -9302,3 +9302,8 @@ bool InvalidTypeSpecializationArity::diagnoseAsError() {
                                   /*generic=*/nullptr);
   return true;
 }
+
+bool InvalidTypeAsKeyPathSubscriptIndex::diagnoseAsError() {
+  emitDiagnostic(diag::cannot_convert_type_to_keypath_subscript_index, ArgType);
+  return true;
+}

lib/Sema/CSDiagnostics.h

@@ -3128,6 +3128,24 @@ class InvalidTypeSpecializationArity final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+/// Diagnose attempts to pass non-keypath as an argument to key path subscript:
+///
+/// \code
+/// func test(data: Int) {
+///   data[keyPath: 42] // error `42` is not a key path
+/// }
+/// \endcode
+class InvalidTypeAsKeyPathSubscriptIndex final : public FailureDiagnostic {
+  Type ArgType;
+
+public:
+  InvalidTypeAsKeyPathSubscriptIndex(const Solution &solution, Type argType,
+                                     ConstraintLocator *locator)
+      : FailureDiagnostic(solution, locator), ArgType(resolveType(argType)) {}
+
+  bool diagnoseAsError() override;
+};
+
 } // end namespace constraints
 } // end namespace swift
 

lib/Sema/CSFix.cpp

@@ -2627,3 +2627,16 @@ IgnoreGenericSpecializationArityMismatch::create(ConstraintSystem &cs,
   return new (cs.getAllocator()) IgnoreGenericSpecializationArityMismatch(
       cs, decl, numParams, numArgs, hasParameterPack, locator);
 }
+
+bool IgnoreKeyPathSubscriptIndexMismatch::diagnose(const Solution &solution,
+                                                   bool asNote) const {
+  InvalidTypeAsKeyPathSubscriptIndex failure(solution, ArgType, getLocator());
+  return failure.diagnose(asNote);
+}
+
+IgnoreKeyPathSubscriptIndexMismatch *
+IgnoreKeyPathSubscriptIndexMismatch::create(ConstraintSystem &cs, Type argType,
+                                            ConstraintLocator *locator) {
+  return new (cs.getAllocator())
+      IgnoreKeyPathSubscriptIndexMismatch(cs, argType, locator);
+}

lib/Sema/CSSimplify.cpp

@@ -12414,9 +12414,18 @@ ConstraintSystem::simplifyKeyPathApplicationConstraint(
     // Otherwise, we don't have a key path type at all.
     return SolutionKind::Error;
   }
-  if (!keyPathTy->isTypeVariableOrMember())
+
+  if (!keyPathTy->isTypeVariableOrMember()) {
+    if (shouldAttemptFixes()) {
+      auto *fix = IgnoreKeyPathSubscriptIndexMismatch::create(
+          *this, keyPathTy, getConstraintLocator(locator));
+      recordAnyTypeVarAsPotentialHole(valueTy);
+      return recordFix(fix) ? SolutionKind::Error : SolutionKind::Solved;
+    }
+
     return SolutionKind::Error;
-  
+  }
+
   return unsolved();
 }
 
@@ -14886,7 +14895,8 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AllowAssociatedValueMismatch:
   case FixKind::GenericArgumentsMismatch:
   case FixKind::AllowConcreteTypeSpecialization:
-  case FixKind::IgnoreGenericSpecializationArityMismatch: {
+  case FixKind::IgnoreGenericSpecializationArityMismatch:
+  case FixKind::IgnoreKeyPathSubscriptIndexMismatch: {
     return recordFix(fix) ? SolutionKind::Error : SolutionKind::Solved;
   }
   case FixKind::IgnoreThrownErrorMismatch: {

lib/Sema/ConstraintLocator.cpp

@@ -111,6 +111,7 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   case ConstraintLocator::PackExpansionType:
   case ConstraintLocator::ThrownErrorType:
   case ConstraintLocator::FallbackType:
+  case ConstraintLocator::KeyPathSubscriptIndex:
     return 0;
 
   case ConstraintLocator::FunctionArgument:
@@ -528,6 +529,10 @@ void LocatorPathElt::dump(raw_ostream &out) const {
   case ConstraintLocator::FallbackType: {
     out << "fallback type";
     break;
+
+  case ConstraintLocator::KeyPathSubscriptIndex:
+    out << "key path subscript index parameter";
+    break;
   }
   }
 }

lib/Sema/ConstraintSystem.cpp

@@ -3693,16 +3693,13 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     // Key path application looks like a subscript(keyPath: KeyPath<Base, T>).
     // The element type is T or @lvalue T based on the key path subtype and
     // the mutability of the base.
-    auto keyPathIndexTy = createTypeVariable(
-        getConstraintLocator(locator, ConstraintLocator::FunctionArgument),
-        TVO_CanBindToInOut);
+    auto *keyPathIndexLoc =
+        getConstraintLocator(locator, ConstraintLocator::KeyPathSubscriptIndex);
+    auto keyPathIndexTy = createTypeVariable(keyPathIndexLoc,
+                                             /*options=*/0);
     auto elementTy = createTypeVariable(
-            getConstraintLocator(locator, ConstraintLocator::FunctionArgument),
-            TVO_CanBindToLValue | TVO_CanBindToNoEscape);
-    auto elementObjTy = createTypeVariable(
-        getConstraintLocator(locator, ConstraintLocator::FunctionArgument),
-        TVO_CanBindToNoEscape);
-    addConstraint(ConstraintKind::Equal, elementTy, elementObjTy, locator);
+        getConstraintLocator(keyPathIndexLoc, ConstraintLocator::KeyPathValue),
+        TVO_CanBindToLValue | TVO_CanBindToNoEscape);
 
     // The element result is an lvalue or rvalue based on the key path class.
     addKeyPathApplicationConstraint(
@@ -6104,6 +6101,7 @@ void constraints::simplifyLocator(ASTNode &anchor,
     case ConstraintLocator::OptionalPayload:
     case ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice:
     case ConstraintLocator::FallbackType:
+    case ConstraintLocator::KeyPathSubscriptIndex:
       break;
     }
 

lib/Sema/TypeCheckConstraints.cpp

@@ -152,6 +152,11 @@ bool TypeVariableType::Implementation::isKeyPathValue() const {
   return locator && locator->isKeyPathValue();
 }
 
+bool TypeVariableType::Implementation::isKeyPathSubscriptIndex() const {
+  return locator &&
+         locator->isLastElement<LocatorPathElt::KeyPathSubscriptIndex>();
+}
+
 bool TypeVariableType::Implementation::isSubscriptResultType() const {
   if (!(locator && locator->getAnchor()))
     return false;

test/Constraints/keypath.swift

@@ -265,3 +265,42 @@ func rdar32101765() {
   let _: KeyPath<R32101765, Float> = \R32101765.prop32101765.unknown
   // expected-error@-1 {{type 'Int' has no member 'unknown'}}
 }
+
+// https://github.com/apple/swift/issues/69795
+func test_invalid_argument_to_keypath_subscript() {
+  func test(x: Int) {
+    x[keyPath: 5]
+    // expected-error@-1 {{cannot use value of type 'Int' as a key path subscript index; argument must be a key path}}
+  }
+
+  let _: (Int) -> Void = {
+    let y = $0
+    y[keyPath: 5]
+    // expected-error@-1 {{cannot use value of type 'Int' as a key path subscript index; argument must be a key path}}
+  }
+
+  func ambiguous(_: (String) -> Void) {}
+  func ambiguous(_: (Int) -> Void) {}
+
+  // FIXME(diagnostic): This is not properly diagnosed in a general case and key path application is even more
+  // complicated because overloads anchored on 'SubscriptExpr -> subscript member' do not point to declarations.
+  // The diagnostic should point out that `ambiguous` is indeed ambiguous and that `5` is not a valid argument
+  // for a key path subscript.
+  ambiguous {
+    // expected-error@-1 {{type of expression is ambiguous without a type annotation}}
+    $0[keyPath: 5]
+  }
+
+  class A {
+  }
+
+  func test_invalid_existential_protocol(base: String, v: any BinaryInteger) {
+    base[keyPath: v]
+    // expected-error@-1 {{cannot use value of type 'any BinaryInteger' as a key path subscript index; argument must be a key path}}
+  }
+
+  func test_invalid_existential_composition(base: String, v: any A & BinaryInteger) {
+    base[keyPath: v]
+    // expected-error@-1 {{cannot use value of type 'A' as a key path subscript index; argument must be a key path}}
+  }
+}

test/expr/unary/keypath/keypath.swift

@@ -1233,3 +1233,7 @@ func test_keypath_coercion_to_function() {
   let fn = \User.email as (User) -> String // Ok
   _ = users.map(fn) // Ok
 }
+
+func test_keypath_application_with_composition(v: String, kp: any KeyPath<String, Int> & PP) {
+  _ = v[keyPath: kp] // Ok
+}