[ConstraintSystem] Initial rework of key path inference

include/swift/AST/DiagnosticsSema.def

@@ -678,9 +678,12 @@ ERROR(expr_swift_keypath_not_starting_with_type,none,
 ERROR(expr_swift_keypath_not_starting_with_dot,none,
       "a Swift key path with contextual root must begin with a leading dot",
       ())
-ERROR(expr_smart_keypath_value_covert_to_contextual_type,none,
+ERROR(expr_keypath_value_covert_to_contextual_type,none,
       "key path value type %0 cannot be converted to contextual type %1",
       (Type, Type))
+ERROR(expr_keypath_type_covert_to_contextual_type,none,
+      "cannot convert key path type %0 to contextual type %1",
+      (Type, Type))
 ERROR(expr_swift_keypath_empty, none,
       "key path must have at least one component", ())
 ERROR(expr_string_interpolation_outside_string,none,

include/swift/Sema/CSBindings.h

@@ -620,10 +620,7 @@ class BindingSet {
 
   void addLiteralRequirement(Constraint *literal);
 
-  void addDefault(Constraint *constraint) {
-    auto defaultTy = constraint->getSecondType();
-    Defaults.insert({defaultTy->getCanonicalType(), constraint});
-  }
+  void addDefault(Constraint *constraint);
 
   /// Check whether the given binding set covers any of the
   /// literal protocols associated with this type variable.

include/swift/Sema/ConstraintLocatorPathElts.def

@@ -281,6 +281,9 @@ SIMPLE_LOCATOR_PATH_ELT(ThrownErrorType)
 /// A type coercion operand.
 SIMPLE_LOCATOR_PATH_ELT(CoercionOperand)
 
+/// A fallback type for some AST location (i.e. key path literal).
+SIMPLE_LOCATOR_PATH_ELT(FallbackType)
+
 #undef LOCATOR_PATH_ELT
 #undef CUSTOM_LOCATOR_PATH_ELT
 #undef SIMPLE_LOCATOR_PATH_ELT

include/swift/Sema/ConstraintSystem.h

@@ -339,6 +339,12 @@ enum TypeVariableOptions {
   TVO_PackExpansion = 0x40,
 };
 
+enum class KeyPathCapability : uint8_t {
+  ReadOnly,
+  Writable,
+  ReferenceWritable
+};
+
 /// The implementation object for a type variable used within the
 /// constraint-solving type checker.
 ///
@@ -5513,6 +5519,25 @@ class ConstraintSystem {
   bool isMemberAvailableOnExistential(Type baseTy,
                                       const ValueDecl *member) const;
 
+  /// Attempts to infer a capability of a key path (i.e. whether it
+  /// is read-only, writable, etc.) based on the referenced members.
+  ///
+  /// \param keyPath The key path literal expression.
+  ///
+  /// \returns `bool` to indicate whether key path is valid or not,
+  /// and capability if it could be determined.
+  std::pair</*isValid=*/bool, llvm::Optional<KeyPathCapability>>
+  inferKeyPathLiteralCapability(KeyPathExpr *keyPath);
+
+  /// A convenience overload of \c inferKeyPathLiteralCapability.
+  ///
+  /// \param keyPathType The type variable that represents the key path literal.
+  ///
+  /// \returns `bool` to indicate whether key path is valid or not,
+  /// and capability if it could be determined.
+  std::pair</*isValid=*/bool, llvm::Optional<KeyPathCapability>>
+  inferKeyPathLiteralCapability(TypeVariableType *keyPathType);
+
   SWIFT_DEBUG_DUMP;
   SWIFT_DEBUG_DUMPER(dump(Expr *));
 

lib/Sema/CSBindings.cpp

@@ -94,6 +94,11 @@ 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())
+    return true;
+
   if (isHole()) {
     auto *locator = TypeVar->getImpl().getLocator();
     assert(locator && "a hole without locator?");
@@ -169,6 +174,12 @@ bool BindingSet::isPotentiallyIncomplete() const {
   if (Info.isGenericParameter())
     return true;
 
+  // Key path literal type is incomplete until there is a
+  // contextual type or key path is resolved enough to infer
+  // capability and promote default into a binding.
+  if (TypeVar->getImpl().isKeyPathType())
+    return Bindings.empty();
+
   // If current type variable is associated with a code completion token
   // it's possible that it doesn't have enough contextual information
   // to be resolved to anything so let's delay considering it until everything
@@ -873,6 +884,65 @@ void PotentialBindings::addDefault(Constraint *constraint) {
   Defaults.insert(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});
+}
+
 bool LiteralRequirement::isCoveredBy(Type type, ConstraintSystem &CS) const {
   auto coversDefaultType = [](Type type, Type defaultType) -> bool {
     if (!defaultType->hasUnboundGenericType())
@@ -2097,6 +2167,10 @@ bool TypeVarBindingProducer::computeNext() {
   }
 
   if (newBindings.empty()) {
+    // If key path type had contextual types, let's not attempt fallback.
+    if (TypeVar->getImpl().isKeyPathType() && !ExploredTypes.empty())
+      return false;
+
     // Add defaultable constraints (if any).
     for (auto *constraint : DelayedDefaults) {
       if (constraint->getKind() == ConstraintKind::FallbackType) {

lib/Sema/CSDiagnostics.cpp

@@ -2497,7 +2497,7 @@ bool ContextualFailure::diagnoseAsError() {
   if (path.empty()) {
     if (auto *KPE = getAsExpr<KeyPathExpr>(anchor)) {
       emitDiagnosticAt(KPE->getLoc(),
-                       diag::expr_smart_keypath_value_covert_to_contextual_type,
+                       diag::expr_keypath_type_covert_to_contextual_type,
                        getFromType(), getToType());
       return true;
     }
@@ -2744,6 +2744,11 @@ bool ContextualFailure::diagnoseAsError() {
     break;
   }
 
+  case ConstraintLocator::KeyPathValue: {
+    diagnostic = diag::expr_keypath_value_covert_to_contextual_type;
+    break;
+  }
+
   default:
     return false;
   }
@@ -7646,7 +7651,7 @@ bool ArgumentMismatchFailure::diagnoseKeyPathAsFunctionResultMismatch() const {
         paramFnType->getParams().front().getPlainType()->isEqual(kpRootType)))
     return false;
 
-  emitDiagnostic(diag::expr_smart_keypath_value_covert_to_contextual_type,
+  emitDiagnostic(diag::expr_keypath_value_covert_to_contextual_type,
                  kpValueType, paramFnType->getResult());
   return true;
 }

lib/Sema/CSGen.cpp

@@ -3742,9 +3742,10 @@ namespace {
         // Allow \Derived.property to be inferred as \Base.property to
         // simulate a sort of covariant conversion from
         // KeyPath<Derived, T> to KeyPath<Base, T>.
-        CS.addConstraint(ConstraintKind::Subtype, rootObjectTy, root, locator);
+        CS.addConstraint(ConstraintKind::Subtype, rootObjectTy, root,
+                         rootLocator);
       }
-      
+
       bool didOptionalChain = false;
       // We start optimistically from an lvalue base.
       Type base = LValueType::get(root);
@@ -3866,6 +3867,12 @@ namespace {
         base = optTy;
       }
 
+      // If we have a malformed KeyPathExpr e.g. let _: KeyPath<A, C> = \A
+      // let's record a AllowKeyPathMissingComponent fix.
+      if (E->hasSingleInvalidComponent()) {
+        (void)CS.recordFix(AllowKeyPathWithoutComponents::create(CS, locator));
+      }
+
       auto valueLocator =
           CS.getConstraintLocator(E, ConstraintLocator::KeyPathValue);
       auto *value = CS.createTypeVariable(valueLocator, TVO_CanBindToNoEscape |
@@ -3884,6 +3891,13 @@ namespace {
 
       CS.addKeyPathConstraint(kpTy, root, value, componentTypeVars, locator);
 
+      // Add a fallback constraint so we have an anchor to use for
+      // capability inference when there is no context.
+      CS.addUnsolvedConstraint(Constraint::create(
+          CS, ConstraintKind::FallbackType, kpTy,
+          BoundGenericType::get(kpDecl, /*parent=*/Type(), {root, value}),
+          CS.getConstraintLocator(E, ConstraintLocator::FallbackType)));
+
       return kpTy;
     }