Merge pull request #84479 from tshortli/expr-availability-key-path-mutability-in-conversions

Sema: Infer key path immutability during availability checking

Author: tshortli

include/swift/AST/Types.h

@@ -1103,6 +1103,12 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Check if this is a std.string type from C++.
   bool isCxxString();
 
+  /// Check if this type is known to represent key paths.
+  bool isKnownKeyPathType();
+
+  /// Check if this type is known to represent immutable key paths.
+  bool isKnownImmutableKeyPathType();
+
   /// Check if this is either an Array, Set or Dictionary collection type defined
   /// at the top level of the Swift module
   bool isKnownStdlibCollectionType();

include/swift/Sema/ConstraintSystem.h

@@ -6525,10 +6525,6 @@ class TypeVarRefCollector : public ASTWalker {
   }
 };
 
-/// Determine whether given type is a known one
-/// for a key path `{Any, Partial, Writable, ReferenceWritable}KeyPath`.
-bool isKnownKeyPathType(Type type);
-
 /// Determine whether the given type is a PartialKeyPath and
 /// AnyKeyPath or existential type thererof, for example,
 /// `PartialKeyPath<...> & Sendable`.

lib/AST/Type.cpp

@@ -1309,6 +1309,15 @@ bool TypeBase::isCxxString() {
          ctx.Id_basic_string.is(clangDecl->getName());
 }
 
+bool TypeBase::isKnownKeyPathType() {
+  return isKeyPath() || isWritableKeyPath() || isReferenceWritableKeyPath() ||
+         isPartialKeyPath() || isAnyKeyPath();
+}
+
+bool TypeBase::isKnownImmutableKeyPathType() {
+  return isKeyPath() || isPartialKeyPath() || isAnyKeyPath();
+}
+
 bool TypeBase::isKnownStdlibCollectionType() {
   if (isArray() || isDictionary() || isSet()) {
     return true;

lib/IDE/SelectedOverloadInfo.cpp

@@ -79,7 +79,7 @@ swift::ide::getSelectedOverloadInfo(const Solution &S,
         fnType->getParams()[0].getPlainType()->castTo<BoundGenericType>();
 
     auto *keyPathDecl = keyPathTy->getAnyNominal();
-    assert(isKnownKeyPathType(keyPathTy) &&
+    assert(keyPathTy->isKnownKeyPathType() &&
            "parameter is supposed to be a keypath");
 
     auto KeyPathDynamicLocator = CS.getConstraintLocator(

lib/Sema/CSApply.cpp

@@ -2689,7 +2689,7 @@ namespace {
       Type keyPathTy = paramType;
       if (auto *existential = keyPathTy->getAs<ExistentialType>()) {
         keyPathTy = existential->getSuperclass();
-        assert(isKnownKeyPathType(keyPathTy));
+        assert(keyPathTy->isKnownKeyPathType());
       }
 
       SmallVector<Component, 2> components;

lib/Sema/CSBindings.cpp

@@ -513,7 +513,7 @@ void BindingSet::inferTransitiveBindings() {
           auto bindingTy = binding.BindingType->lookThroughAllOptionalTypes();
 
           Type inferredRootTy;
-          if (isKnownKeyPathType(bindingTy)) {
+          if (bindingTy->isKnownKeyPathType()) {
             // AnyKeyPath doesn't have a root type.
             if (bindingTy->isAnyKeyPath())
               continue;
@@ -728,7 +728,8 @@ bool BindingSet::finalize(bool transitive) {
       for (const auto &binding : Bindings) {
         auto bindingTy = binding.BindingType->lookThroughAllOptionalTypes();
 
-        assert(isKnownKeyPathType(bindingTy) || bindingTy->is<FunctionType>());
+        assert(bindingTy->isKnownKeyPathType() ||
+               bindingTy->is<FunctionType>());
 
         // Functions don't have capability so we can simply add them.
         if (auto *fnType = bindingTy->getAs<FunctionType>()) {
@@ -1706,15 +1707,15 @@ PotentialBindings::inferFromRelational(ConstraintSystem &CS,
       if (type->isExistentialType()) {
         auto layout = type->getExistentialLayout();
         if (auto superclass = layout.explicitSuperclass) {
-          if (isKnownKeyPathType(superclass)) {
+          if (superclass->isKnownKeyPathType()) {
             type = superclass;
             objectTy = superclass;
           }
         }
       }
     }
 
-    if (!(isKnownKeyPathType(objectTy) || objectTy->is<AnyFunctionType>()))
+    if (!(objectTy->isKnownKeyPathType() || objectTy->is<AnyFunctionType>()))
       return std::nullopt;
   }
 

lib/Sema/CSDiagnostics.cpp

@@ -2262,7 +2262,7 @@ bool AssignmentFailure::diagnoseAsError() {
 
       auto type = getType(immutableExpr);
 
-      if (isKnownKeyPathType(type))
+      if (type->isKnownKeyPathType())
         message += " is read-only";
       else if (VD->isCaptureList())
         message += " is an immutable capture";
@@ -7487,7 +7487,7 @@ bool ArgumentMismatchFailure::diagnoseAsError() {
 
   // Unresolved key path argument requires a tailored diagnostic
   // that doesn't mention a fallback type - `KeyPath<_, _>`.
-  if (argType->isKeyPath() && !isKnownKeyPathType(paramType)) {
+  if (argType->isKeyPath() && !paramType->isKnownKeyPathType()) {
     auto keyPathTy = argType->castTo<BoundGenericType>();
     auto rootTy = keyPathTy->getGenericArgs()[0];
     if (rootTy->is<UnresolvedType>()) {
@@ -7852,7 +7852,7 @@ bool ArgumentMismatchFailure::diagnoseKeyPathAsFunctionResultMismatch() const {
   auto argType = getFromType();
   auto paramType = getToType();
 
-  if (!isKnownKeyPathType(argType))
+  if (!argType->isKnownKeyPathType())
     return false;
 
   auto kpType = argType->castTo<BoundGenericType>();

lib/Sema/CSSimplify.cpp

@@ -5598,7 +5598,7 @@ bool ConstraintSystem::repairFailures(
     // fix-up here unless last component has already a invalid type or
     // instance fix recorded.
     if (isExpr<KeyPathExpr>(anchor)) {
-      if (isKnownKeyPathType(lhs) && isKnownKeyPathType(rhs)) {
+      if (lhs->isKnownKeyPathType() && rhs->isKnownKeyPathType()) {
         // If we have a conversion happening here, we should let fix happen in
         // simplifyRestrictedConstraint.
         if (hasAnyRestriction())

lib/Sema/ConstraintSystem.cpp

@@ -4363,12 +4363,6 @@ Solution::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
                                    fnInterfaceType, fnType, callee);
 }
 
-bool constraints::isKnownKeyPathType(Type type) {
-  return type->isKeyPath() || type->isWritableKeyPath() ||
-         type->isReferenceWritableKeyPath() || type->isPartialKeyPath() ||
-         type->isAnyKeyPath();
-}
-
 bool constraints::isTypeErasedKeyPathType(Type type) {
   assert(type);
 

lib/Sema/TypeCheckAvailability.cpp

@@ -29,6 +29,7 @@
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/DeclExportabilityVisitor.h"
 #include "swift/AST/DiagnosticsParse.h"
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/NameLookup.h"
@@ -2232,9 +2233,6 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     /// The context does not involve a key path access.
     None,
 
-    /// The context is an expression that is coerced to a read-only key path.
-    ReadOnlyCoercion,
-
     /// The context is a key path application (`x[keyPath: \.member]`).
     Application,
   };
@@ -2361,15 +2359,6 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
                                           FCE->getLoc(),
                                           Where.getDeclContext());
     }
-    if (auto DTBE = dyn_cast<DerivedToBaseExpr>(E)) {
-      if (auto ty = DTBE->getType()) {
-        if (ty->isKeyPath()) {
-          walkInKeyPathAccessContext(DTBE->getSubExpr(),
-                                     KeyPathAccessContext::ReadOnlyCoercion);
-          return Action::SkipChildren(E);
-        }
-      }
-    }
     if (auto KP = dyn_cast<KeyPathExpr>(E)) {
       maybeDiagKeyPath(KP);
     }
@@ -2586,19 +2575,49 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
 
   StorageAccessKind getStorageAccessKindForKeyPath(KeyPathExpr *KP) const {
     switch (KeyPathAccess) {
-    case KeyPathAccessContext::None:
-      // Use the key path's type to determine the access kind.
-      if (!KP->isObjC())
-        if (auto keyPathType = KP->getKeyPathType())
-          if (keyPathType->isKeyPath())
+    case KeyPathAccessContext::None: {
+      // Obj-C key paths are always considered mutable.
+      if (KP->isObjC())
+        return StorageAccessKind::GetSet;
+
+      // Check whether key path type indicates immutability, in which case only
+      // the getter will be used.
+      if (auto keyPathType = KP->getKeyPathType())
+        if (keyPathType->isKnownImmutableKeyPathType())
+          return StorageAccessKind::Get;
+
+      // Check if there's a conversion expression containing this one directly
+      // above it in the stack. If so, and that conversion is to an immutable
+      // key path type, then we should infer that use of the key path only
+      // implies use of the getter.
+      ArrayRef<const Expr *> exprs = ExprStack;
+      // Must be called while visiting the given key path expression.
+      ASSERT(exprs.back() == KP);
+
+      for (auto *expr : llvm::reverse(exprs.drop_back())) {
+        // Only traverse through conversions on the stack.
+        if (!isa<ImplicitConversionExpr>(expr) &&
+            !isa<OpenExistentialExpr>(expr))
+          break;
+
+        if (auto ty = expr->getType()) {
+          if (ty->isKnownImmutableKeyPathType())
             return StorageAccessKind::Get;
 
-      return StorageAccessKind::GetSet;
+          if (auto existential = dyn_cast<ExistentialType>(ty)) {
+            if (auto superclass =
+                    existential->getExistentialLayout().getSuperclass()) {
+              if (superclass->isKnownImmutableKeyPathType())
+                return StorageAccessKind::Get;
+            }
+          }
+        }
+      }
 
-    case KeyPathAccessContext::ReadOnlyCoercion:
-      // The type of this key path is being coerced to the type KeyPath<_, _> so
-      // ignore the actual key path type.
-      return StorageAccessKind::Get;
+      // We failed to prove that the key path is only used immutably,
+      // so diagnose both get and set access.
+      return StorageAccessKind::GetSet;
+    }
 
     case KeyPathAccessContext::Application:
       // The key path is being applied directly to a base so treat it as if it