[CSBindings] Delay key path type inference until literal capability is known

Since key path root is now transitively inferred. Key path type
inference can be delayed until key path is resolved enough to
infer its capability.

This solves multiple problems:

- Inference fully controls what key path type is bound to;
- KeyPath constraint simplification doesn't have to double-check
  the capability and attempt to re-bind key path type;
- Custom logic to resolve key path type is no longer necessary;
- Diagnostics are improved because capability and root/value type
  mismatch are diagnosed when key path is matched against the
  contextual type.

Author: xedin

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;
 

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
@@ -531,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);
@@ -573,6 +591,90 @@ 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);
+
+      if (!isValid) {
+        // If key path is invalid we have to drop all the contextual
+        // bindings, none of the could be used unless capability is
+        // known.
+        Bindings.clear();
+
+        // 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()) {
+          auto rootTy = CS.getKeyPathRootType(keyPath);
+          // 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());
+            addBinding(
+                {fallback->first, AllowedBindingKind::Exact, fallback->second});
+          } else {
+            addBinding(PotentialBinding::forHole(
+                TypeVar, CS.getConstraintLocator(
+                             keyPath, ConstraintLocator::FallbackType)));
+          }
+        }
+
+        // No need for fallback if key path is invalid.
+        Defaults.clear();
+        return;
+      }
+
+      // If the key path is sufficiently resolved we can add inferred binding
+      // to the set.
+      if (capability) {
+        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 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, CS.getKeyPathRootType(keyPath),
+                           CS.getKeyPathValueType(keyPath));
+
+        updatedBindings.insert(
+            {keyPathTy, AllowedBindingKind::Exact, keyPathLoc});
+
+        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
@@ -927,60 +1029,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/CSSimplify.cpp

@@ -4471,13 +4471,6 @@ 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));
 

test/expr/unary/keypath/keypath.swift

@@ -692,7 +692,8 @@ func testSubtypeKeypathClass(_ keyPath: ReferenceWritableKeyPath<Base, Int>) {
 
 func testSubtypeKeypathProtocol(_ keyPath: ReferenceWritableKeyPath<PP, Int>) {
   testSubtypeKeypathProtocol(\Base.i)
-  // expected-error@-1 {{key path with root type 'any PP' cannot be applied to a base of type 'Base'}}
+  // expected-error@-1 {{cannot convert value of type 'ReferenceWritableKeyPath<Base, Int>' to expected argument type 'ReferenceWritableKeyPath<any PP, Int>'}}
+  // expected-note@-2 {{arguments to generic parameter 'Root' ('Base' and 'any PP') are expected to be equal}}
 }
 
 // rdar://problem/32057712