Merge pull request #69742 from xedin/transitive-keypath-base-inference

[ConstraintSystem] More changes to key path inference

Author: xedin

include/swift/AST/DiagnosticsSema.def

@@ -376,6 +376,10 @@ ERROR(cannot_convert_argument_value,none,
       "cannot convert value of type %0 to expected argument type %1",
       (Type,Type))
 
+ERROR(cannot_convert_unresolved_key_path_argument_value,none,
+      "cannot convert value of key path type to expected argument type %0",
+      (Type))
+
 ERROR(cannot_convert_argument_value_for_swift_func,none,
       "cannot convert value of type %0 to expected argument type %1 "
       "because %kind2 was not imported from C header",

include/swift/Sema/CSBindings.h

@@ -71,6 +71,8 @@ enum class LiteralBindingKind : uint8_t {
 /// along with information that can be used to construct related
 /// bindings, e.g., the supertypes of a given type.
 struct PotentialBinding {
+  friend class BindingSet;
+
   /// The type to which the type variable can be bound.
   Type BindingType;
 

include/swift/Sema/ConstraintSystem.h

@@ -491,6 +491,10 @@ class TypeVariableType::Implementation {
   /// a type of a key path expression.
   bool isKeyPathType() const;
 
+  /// Determine whether this type variable represents a root type of a key path
+  /// expression.
+  bool isKeyPathRoot() const;
+
   /// Determine whether this type variable represents a value type of a key path
   /// expression.
   bool isKeyPathValue() const;
@@ -4008,20 +4012,6 @@ 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/CSBindings.cpp

@@ -96,7 +96,7 @@ bool BindingSet::isDelayed() const {
 
   // Delay key path literal type binding until there is at least
   // one contextual binding (or default is promoted into a binding).
-  if (TypeVar->getImpl().isKeyPathType() && Bindings.empty())
+  if (TypeVar->getImpl().isKeyPathType() && !Defaults.empty())
     return true;
 
   if (isHole()) {
@@ -178,7 +178,7 @@ bool BindingSet::isPotentiallyIncomplete() const {
   // contextual type or key path is resolved enough to infer
   // capability and promote default into a binding.
   if (TypeVar->getImpl().isKeyPathType())
-    return Bindings.empty();
+    return !Defaults.empty();
 
   // If current type variable is associated with a code completion token
   // it's possible that it doesn't have enough contextual information
@@ -430,6 +430,41 @@ void BindingSet::inferTransitiveBindings(
         &inferredBindings) {
   using BindingKind = AllowedBindingKind;
 
+  // If the current type variable represents a key path root type
+  // let's try to transitively infer its type through bindings of
+  // a key path type.
+  if (TypeVar->getImpl().isKeyPathRoot()) {
+    auto *locator = TypeVar->getImpl().getLocator();
+    if (auto *keyPathTy =
+            CS.getType(locator->getAnchor())->getAs<TypeVariableType>()) {
+      auto keyPathBindings = inferredBindings.find(keyPathTy);
+      if (keyPathBindings != inferredBindings.end()) {
+        auto &bindings = keyPathBindings->getSecond();
+
+        for (auto &binding : bindings.Bindings) {
+          auto bindingTy = binding.BindingType->lookThroughAllOptionalTypes();
+
+          Type inferredRootTy;
+          if (isKnownKeyPathType(bindingTy)) {
+            // AnyKeyPath doesn't have a root type.
+            if (bindingTy->isAnyKeyPath())
+              continue;
+
+            auto *BGT = bindingTy->castTo<BoundGenericType>();
+            inferredRootTy = BGT->getGenericArgs()[0];
+          } else if (auto *fnType = bindingTy->getAs<FunctionType>()) {
+            if (fnType->getNumParams() == 1)
+              inferredRootTy = fnType->getParams()[0].getParameterType();
+          }
+
+          if (inferredRootTy && !inferredRootTy->isTypeVariableOrMember())
+            addBinding(
+                binding.withSameSource(inferredRootTy, BindingKind::Exact));
+        }
+      }
+    }
+  }
+
   for (const auto &entry : Info.SupertypeOf) {
     auto relatedBindings = inferredBindings.find(entry.first);
     if (relatedBindings == inferredBindings.end())
@@ -496,6 +531,24 @@ void BindingSet::inferTransitiveBindings(
   }
 }
 
+static BoundGenericType *getKeyPathType(ASTContext &ctx,
+                                        KeyPathCapability capability,
+                                        Type rootType, Type valueType) {
+  switch (capability) {
+  case KeyPathCapability::ReadOnly:
+    return BoundGenericType::get(ctx.getKeyPathDecl(), /*parent=*/Type(),
+                                 {rootType, valueType});
+
+  case KeyPathCapability::Writable:
+    return BoundGenericType::get(ctx.getWritableKeyPathDecl(),
+                                 /*parent=*/Type(), {rootType, valueType});
+
+  case KeyPathCapability::ReferenceWritable:
+    return BoundGenericType::get(ctx.getReferenceWritableKeyPathDecl(),
+                                 /*parent=*/Type(), {rootType, valueType});
+  }
+}
+
 void BindingSet::finalize(
     llvm::SmallDenseMap<TypeVariableType *, BindingSet> &inferredBindings) {
   inferTransitiveBindings(inferredBindings);
@@ -538,6 +591,84 @@ void BindingSet::finalize(
       }
     }
 
+    if (TypeVar->getImpl().isKeyPathType()) {
+      auto &ctx = CS.getASTContext();
+
+      auto *keyPathLoc = TypeVar->getImpl().getLocator();
+      auto *keyPath = castToExpr<KeyPathExpr>(keyPathLoc->getAnchor());
+
+      bool isValid;
+      llvm::Optional<KeyPathCapability> capability;
+
+      std::tie(isValid, capability) = CS.inferKeyPathLiteralCapability(TypeVar);
+
+      // Key path literal is not yet sufficiently resolved.
+      if (isValid && !capability)
+        return;
+
+      // If the key path is sufficiently resolved we can add inferred binding
+      // to the set.
+      SmallSetVector<PotentialBinding, 4> updatedBindings;
+      for (const auto &binding : Bindings) {
+        auto bindingTy = binding.BindingType->lookThroughAllOptionalTypes();
+
+        assert(isKnownKeyPathType(bindingTy) || bindingTy->is<FunctionType>());
+
+        // Functions don't have capability so we can simply add them.
+        if (bindingTy->is<FunctionType>())
+          updatedBindings.insert(binding);
+      }
+
+      // Note that even though key path literal maybe be invalid it's
+      // still the best course of action to use contextual function type
+      // bindings because they allow to propagate type information from
+      // the key path into the context, so key path bindings are addded
+      // only if there is absolutely no other choice.
+      if (updatedBindings.empty()) {
+        auto rootTy = CS.getKeyPathRootType(keyPath);
+
+        // A valid key path literal.
+        if (capability) {
+          // Note that the binding is formed using root & value
+          // type variables produced during constraint generation
+          // because at this point root is already known (otherwise
+          // inference wouldn't been able to determine key path's
+          // capability) and we always want to infer value from
+          // the key path and match it to a contextual type to produce
+          // better diagnostics.
+          auto keyPathTy = getKeyPathType(ctx, *capability, rootTy,
+                                          CS.getKeyPathValueType(keyPath));
+          updatedBindings.insert(
+              {keyPathTy, AllowedBindingKind::Exact, keyPathLoc});
+        } else if (CS.shouldAttemptFixes()) {
+          // If key path is structurally correct and has a resolved root
+          // type, let's promote the fallback type into a binding because
+          // root would have been inferred from explicit type already and
+          // it's benefitial for diagnostics to assign a non-placeholder
+          // type to key path literal to propagate root/value to the context.
+          if (!keyPath->hasSingleInvalidComponent() &&
+              (keyPath->getParsedRoot() ||
+               !CS.getFixedType(rootTy)->isTypeVariableOrMember())) {
+            auto fallback = llvm::find_if(Defaults, [](const auto &entry) {
+              return entry.second->getKind() == ConstraintKind::FallbackType;
+            });
+            assert(fallback != Defaults.end());
+            updatedBindings.insert(
+                {fallback->first, AllowedBindingKind::Exact, fallback->second});
+          } else {
+            updatedBindings.insert(PotentialBinding::forHole(
+                TypeVar, CS.getConstraintLocator(
+                             keyPath, ConstraintLocator::FallbackType)));
+          }
+        }
+      }
+
+      Bindings = std::move(updatedBindings);
+      Defaults.clear();
+
+      return;
+    }
+
     if (CS.shouldAttemptFixes() &&
         locator->isLastElement<LocatorPathElt::UnresolvedMemberChainResult>()) {
       // Let's see whether this chain is valid, if it isn't then to avoid
@@ -795,6 +926,12 @@ llvm::Optional<BindingSet> ConstraintSystem::determineBestBindings(
   auto isViableForRanking = [this](const BindingSet &bindings) -> bool {
     auto *typeVar = bindings.getTypeVariable();
 
+    // Key path root type variable is always viable because it can be
+    // transitively inferred from key path type during binding set
+    // finalization.
+    if (typeVar->getImpl().isKeyPathRoot())
+      return true;
+
     // Type variable representing a base of unresolved member chain should
     // always be considered viable for ranking since it's allow to infer
     // types from transitive protocol requirements.
@@ -886,60 +1023,6 @@ void PotentialBindings::addDefault(Constraint *constraint) {
 
 void BindingSet::addDefault(Constraint *constraint) {
   auto defaultTy = constraint->getSecondType();
-
-  if (TypeVar->getImpl().isKeyPathType() && Bindings.empty()) {
-    if (constraint->getKind() == ConstraintKind::FallbackType) {
-      auto &ctx = CS.getASTContext();
-
-      bool isValid;
-      llvm::Optional<KeyPathCapability> capability;
-
-      std::tie(isValid, capability) = CS.inferKeyPathLiteralCapability(TypeVar);
-
-      if (!isValid) {
-        // If one of the references in a key path is invalid let's add
-        // a placeholder binding in diagnostic mode to indicate that
-        // the key path cannot be properly resolved.
-        if (CS.shouldAttemptFixes()) {
-          addBinding({PlaceholderType::get(ctx, TypeVar),
-                      AllowedBindingKind::Exact, constraint});
-        }
-
-        // During normal solving the set has to stay empty.
-        return;
-      }
-
-      if (capability) {
-        auto *keyPathType = defaultTy->castTo<BoundGenericType>();
-
-        auto root = keyPathType->getGenericArgs()[0];
-        auto value = keyPathType->getGenericArgs()[1];
-
-        switch (*capability) {
-        case KeyPathCapability::ReadOnly:
-          break;
-
-        case KeyPathCapability::Writable:
-          keyPathType = BoundGenericType::get(ctx.getWritableKeyPathDecl(),
-                                              /*parent=*/Type(), {root, value});
-          break;
-
-        case KeyPathCapability::ReferenceWritable:
-          keyPathType =
-              BoundGenericType::get(ctx.getReferenceWritableKeyPathDecl(),
-                                    /*parent=*/Type(), {root, value});
-          break;
-        }
-
-        addBinding({keyPathType, AllowedBindingKind::Exact, constraint});
-      }
-
-      // If key path is not yet sufficiently resolved, don't add any
-      // bindings.
-      return;
-    }
-  }
-
   Defaults.insert({defaultTy->getCanonicalType(), constraint});
 }
 

lib/Sema/CSDiagnostics.cpp

@@ -2604,6 +2604,11 @@ bool ContextualFailure::diagnoseAsError() {
     if (diagnoseYieldByReferenceMismatch())
       return true;
 
+    if (isExpr<KeyPathExpr>(anchor)) {
+      diagnostic = diag::expr_keypath_type_covert_to_contextual_type;
+      break;
+    }
+
     if (isExpr<OptionalTryExpr>(anchor) ||
         isExpr<OptionalEvaluationExpr>(anchor)) {
       auto objectType = fromType->getOptionalObjectType();
@@ -7290,6 +7295,18 @@ bool ArgumentMismatchFailure::diagnoseAsError() {
 
   auto argType = getFromType();
 
+  // Unresolved key path argument requires a tailored diagnostic
+  // that doesn't mention a fallback type - `KeyPath<_, _>`.
+  if (argType->isKeyPath() && !isKnownKeyPathType(paramType)) {
+    auto keyPathTy = argType->castTo<BoundGenericType>();
+    auto rootTy = keyPathTy->getGenericArgs()[0];
+    if (rootTy->is<UnresolvedType>()) {
+      emitDiagnostic(diag::cannot_convert_unresolved_key_path_argument_value,
+                     paramType);
+      return true;
+    }
+  }
+
   if (paramType->isAnyObject()) {
     emitDiagnostic(diag::cannot_convert_argument_value_anyobject, argType,
                    paramType);

lib/Sema/CSGen.cpp

@@ -3853,7 +3853,10 @@ namespace {
         // constraint system so we can find it later.
         CS.setType(E, i, base);
       }
-      
+
+      auto valueLocator =
+        CS.getConstraintLocator(E, ConstraintLocator::KeyPathValue);
+
       // If there was an optional chaining component, the end result must be
       // optional.
       if (didOptionalChain) {
@@ -3862,8 +3865,7 @@ namespace {
         componentTypeVars.push_back(objTy);
 
         auto optTy = OptionalType::get(objTy);
-        CS.addConstraint(ConstraintKind::Conversion, base, optTy,
-                         locator);
+        CS.addConstraint(ConstraintKind::Conversion, base, optTy, valueLocator);
         base = optTy;
       }
 
@@ -3873,8 +3875,7 @@ namespace {
         (void)CS.recordFix(AllowKeyPathWithoutComponents::create(CS, locator));
       }
 
-      auto valueLocator =
-          CS.getConstraintLocator(E, ConstraintLocator::KeyPathValue);
+
       auto *value = CS.createTypeVariable(valueLocator, TVO_CanBindToNoEscape |
                                                             TVO_CanBindToHole);
       CS.addConstraint(ConstraintKind::Equal, base, value, valueLocator);