Merge pull request #67297 from amritpan/build-out-resolve-kp

[ConstraintSystem] Set up keypath expressions typechecking

Author: xedin

include/swift/Sema/ConstraintLocator.h

@@ -254,7 +254,7 @@ class ConstraintLocator : public llvm::FoldingSetNode {
 
   /// Determine whether given locator points to the keypath value
   bool isKeyPathValue() const;
-  
+
   /// Determine whether given locator points to the choice picked as
   /// as result of the key path dynamic member lookup operation.
   bool isResultOfKeyPathDynamicMemberLookup() const;

include/swift/Sema/ConstraintSystem.h

@@ -486,6 +486,10 @@ class TypeVariableType::Implementation {
   /// a type of a key path expression.
   bool isKeyPathType() const;
 
+  /// Determine whether this type variable represents a value type of a key path
+  /// expression.
+  bool isKeyPathValue() const;
+
   /// Determine whether this type variable represents a subscript result type.
   bool isSubscriptResultType() const;
 
@@ -3014,6 +3018,20 @@ class ConstraintSystem {
     return nullptr;
   }
 
+  TypeVariableType *getKeyPathValueType(const KeyPathExpr *keyPath) const {
+    auto result = getKeyPathValueTypeIfAvailable(keyPath);
+    assert(result);
+    return result;
+  }
+
+  TypeVariableType *
+  getKeyPathValueTypeIfAvailable(const KeyPathExpr *keyPath) const {
+    auto result = KeyPaths.find(keyPath);
+    if (result != KeyPaths.end())
+      return std::get<1>(result->second);
+    return nullptr;
+  }
+
   TypeBase* getFavoredType(Expr *E) {
     assert(E != nullptr);
     return this->FavoredTypes[E];
@@ -3967,6 +3985,20 @@ class ConstraintSystem {
   bool resolveClosure(TypeVariableType *typeVar, Type contextualType,
                       ConstraintLocatorBuilder locator);
 
+  /// Given the fact a contextual type is now available for the type
+  /// variable representing one of the key path expressions, let's set a
+  /// pre-determined key path expression type.
+  ///
+  /// \param typeVar The type variable representing a key path expression.
+  /// \param contextualType The contextual type this key path would be
+  /// converted to.
+  /// \param locator The locator associated with contextual type.
+  ///
+  /// \returns `true` if it was possible to generate constraints for
+  /// the keyPath expression, `false` otherwise.
+  bool resolveKeyPath(TypeVariableType *typeVar, Type contextualType,
+                      ConstraintLocatorBuilder locator);
+
   /// Given the fact that contextual type is now available for the type
   /// variable representing a pack expansion type, let's resolve the expansion.
   ///

lib/Sema/CSSimplify.cpp

@@ -4439,6 +4439,13 @@ ConstraintSystem::matchTypesBindTypeVar(
                : getTypeMatchFailure(locator);
   }
 
+  if (typeVar->getImpl().isKeyPathType()) {
+    if (flags.contains(TMF_BindingTypeVariable))
+      return resolveKeyPath(typeVar, type, locator)
+                 ? getTypeMatchSuccess()
+                 : getTypeMatchFailure(locator);
+  }
+
   assignFixedType(typeVar, type, /*updateState=*/true,
                   /*notifyInference=*/!flags.contains(TMF_BindingTypeVariable));
 
@@ -4630,6 +4637,13 @@ repairViaBridgingCast(ConstraintSystem &cs, Type fromType, Type toType,
   return true;
 }
 
+/// Return tuple of type and number of optionals on that type.
+static std::pair<Type, unsigned> getObjectTypeAndNumUnwraps(Type type) {
+  SmallVector<Type, 2> optionals;
+  Type objType = type->lookThroughAllOptionalTypes(optionals);
+  return std::make_pair(objType, optionals.size());
+}
+
 static bool
 repairViaOptionalUnwrap(ConstraintSystem &cs, Type fromType, Type toType,
                         ConstraintKind matchKind,
@@ -4747,17 +4761,11 @@ repairViaOptionalUnwrap(ConstraintSystem &cs, Type fromType, Type toType,
     }
   }
 
-  auto getObjectTypeAndUnwraps = [](Type type) -> std::pair<Type, unsigned> {
-    SmallVector<Type, 2> optionals;
-    Type objType = type->lookThroughAllOptionalTypes(optionals);
-    return std::make_pair(objType, optionals.size());
-  };
-
   Type fromObjectType, toObjectType;
   unsigned fromUnwraps, toUnwraps;
 
-  std::tie(fromObjectType, fromUnwraps) = getObjectTypeAndUnwraps(fromType);
-  std::tie(toObjectType, toUnwraps) = getObjectTypeAndUnwraps(toType);
+  std::tie(fromObjectType, fromUnwraps) = getObjectTypeAndNumUnwraps(fromType);
+  std::tie(toObjectType, toUnwraps) = getObjectTypeAndNumUnwraps(toType);
 
   // Since equality is symmetric and it decays into a `Bind`, eagerly
   // unwrapping optionals from either side might be incorrect since
@@ -6491,6 +6499,19 @@ bool ConstraintSystem::repairFailures(
     if (!fromType || !toType)
       break;
 
+    Type fromObjectType, toObjectType;
+    unsigned fromUnwraps, toUnwraps;
+
+    std::tie(fromObjectType, fromUnwraps) = getObjectTypeAndNumUnwraps(lhs);
+    std::tie(toObjectType, toUnwraps) = getObjectTypeAndNumUnwraps(rhs);
+
+    // If the bound contextual type is more optional than the binding type, then
+    // propogate binding type to contextual type and attempt to solve.
+    if (fromUnwraps < toUnwraps) {
+      (void)matchTypes(fromObjectType, toObjectType, ConstraintKind::Bind,
+                       TMF_ApplyingFix, locator);
+    }
+
     // Drop both `GenericType` elements.
     path.pop_back();
     path.pop_back();
@@ -6704,6 +6725,15 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
         return getTypeMatchSuccess();
       }
 
+      // If type variable represents a key path value type, defer binding it to
+      // contextual type in diagnostic mode. We want it to be bound from the
+      // last key path component to help with diagnostics.
+      if (shouldAttemptFixes()) {
+        if (typeVar1 && typeVar1->getImpl().isKeyPathValue() &&
+            !flags.contains(TMF_BindingTypeVariable))
+          return formUnsolvedResult();
+      }
+
       assert((type1->is<TypeVariableType>() != type2->is<TypeVariableType>()) &&
              "Expected a type variable and a non type variable!");
 
@@ -11542,6 +11572,31 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
   return !generateConstraints(AnyFunctionRef{closure}, closure->getBody());
 }
 
+bool ConstraintSystem::resolveKeyPath(TypeVariableType *typeVar,
+                                      Type contextualType,
+                                      ConstraintLocatorBuilder locator) {
+  auto *keyPathLocator = typeVar->getImpl().getLocator();
+  auto *keyPath = castToExpr<KeyPathExpr>(keyPathLocator->getAnchor());
+  if (keyPath->hasSingleInvalidComponent()) {
+    assignFixedType(typeVar, contextualType);
+    return true;
+  }
+  if (auto *BGT = contextualType->getAs<BoundGenericType>()) {
+    auto args = BGT->getGenericArgs();
+    if (isKnownKeyPathType(contextualType) && args.size() >= 1) {
+      auto root = BGT->getGenericArgs()[0];
+
+      auto *keyPathValueTV = getKeyPathValueType(keyPath);
+      contextualType = BoundGenericType::get(
+          args.size() == 1 ? getASTContext().getKeyPathDecl() : BGT->getDecl(),
+          /*parent=*/Type(), {root, keyPathValueTV});
+    }
+  }
+
+  assignFixedType(typeVar, contextualType);
+  return true;
+}
+
 bool ConstraintSystem::resolvePackExpansion(TypeVariableType *typeVar,
                                             Type contextualType) {
   auto *locator = typeVar->getImpl().getLocator();

lib/Sema/TypeCheckConstraints.cpp

@@ -144,6 +144,10 @@ bool TypeVariableType::Implementation::isKeyPathType() const {
   return locator && locator->isKeyPathType();
 }
 
+bool TypeVariableType::Implementation::isKeyPathValue() const {
+  return locator && locator->isKeyPathValue();
+}
+
 bool TypeVariableType::Implementation::isSubscriptResultType() const {
   if (!(locator && locator->getAnchor()))
     return false;

test/Constraints/keypath.swift

@@ -136,8 +136,8 @@ func test_mismatch_with_contextual_optional_result() {
     var arr: [Int] = []
   }
 
-  let _ = A(B(), keyPath: \.arr)
-  // expected-error@-1 {{key path value type '[Int]' cannot be converted to contextual type '[Int]?'}}
+  let _ = A(B(), keyPath: \.arr) // expected-error {{cannot convert value of type 'KeyPath<B, [Int]>' to expected argument type 'KeyPath<B, [Int]?>'}}
+  // expected-note@-1 {{arguments to generic parameter 'Value' ('[Int]' and '[Int]?') are expected to be equal}}
 }
 
 // https://github.com/apple/swift/issues/53581
@@ -180,6 +180,20 @@ func key_path_root_mismatch<T>(_ base: KeyPathBase?, subBase: KeyPathBaseSubtype
 
 }
 
+func key_path_value_mismatch() {
+  struct S {
+    var member: Int
+  }
+	
+  func test(_: KeyPath<S, String>) {}
+  // expected-note@-1 {{found candidate with type 'KeyPath<S, Int>'}}
+  func test(_: KeyPath<S, Float>) {}
+  // expected-note@-1 {{found candidate with type 'KeyPath<S, Int>'}}
+	
+  test(\.member)
+  // expected-error@-1 {{no exact matches in call to local function 'test'}}
+}
+
 // https://github.com/apple/swift/issues/55884
 func f_55884() {
   func f<T>(_ x: KeyPath<String?, T>) -> T { "1"[keyPath: x] }

test/Sema/keypath_subscript_nolabel.swift

@@ -28,4 +28,5 @@ struct S4 {
   subscript(v: KeyPath<S4, String>) -> Int { get { 0 } set(newValue) {} }
 }
 var s4 = S4()
-s4[\.x] = 10 // expected-error {{key path value type 'Int' cannot be converted to contextual type 'String'}}
+s4[\.x] = 10 // expected-error {{cannot convert value of type 'KeyPath<S4, Int>' to expected argument type 'KeyPath<S4, String>'}}
+// expected-note@-1 {{arguments to generic parameter 'Value' ('Int' and 'String') are expected to be equal}}

test/expr/unary/keypath/keypath.swift

@@ -228,7 +228,7 @@ func testNoComponents() {
   let _: KeyPath<A, A> = \A // expected-error{{must have at least one component}}
   let _: KeyPath<C, A> = \C // expected-error{{must have at least one component}}
   // expected-error@-1 {{generic parameter 'T' could not be inferred}}
-  let _: KeyPath<A, C> = \A // expected-error{{must have at least one component}} 
+  let _: KeyPath<A, C> = \A // expected-error{{must have at least one component}}
   // expected-error@-1 {{generic parameter 'T' could not be inferred}}
   _ = \A // expected-error {{key path must have at least one component}}
 }