[CSSimplify] Reject key path if root type is AnyObject (#23820)

Detect situations where `AnyObject` is attempted to be used as a root type of the key path
early and diagnose via new diagnostics framework.

Author: theblixguy

include/swift/AST/DiagnosticsSema.def

@@ -506,6 +506,8 @@ ERROR(expr_keypath_subscript_index_not_hashable, none,
 ERROR(expr_smart_keypath_application_type_mismatch,none,
       "key path of type %0 cannot be applied to a base of type %1",
       (Type, Type))
+ERROR(expr_swift_keypath_anyobject_root,none,
+      "the root type of a Swift key path cannot be 'AnyObject'", ())
 WARNING(expr_deprecated_writable_keypath,none,
         "forming a writable keypath to property %0 that is read-only in this context "
         "is deprecated and will be removed in a future release",(DeclName))

lib/Sema/CSApply.cpp

@@ -4342,15 +4342,6 @@ namespace {
       Type baseTy = keyPathTy->getGenericArgs()[0];
       Type leafTy = keyPathTy->getGenericArgs()[1];
 
-      // We do not allow keypaths to go through AnyObject
-      if (baseTy->isAnyObject()) {
-        auto rootTyRepr = E->getRootType();
-        cs.TC.diagnose(rootTyRepr->getLoc(),
-                       diag::expr_swift_keypath_invalid_component)
-             .highlight(rootTyRepr->getSourceRange());
-        return nullptr;
-      }
-
       for (unsigned i : indices(E->getComponents())) {
         auto &origComponent = E->getMutableComponents()[i];
 

lib/Sema/CSDiag.cpp

@@ -758,6 +758,7 @@ void FailureDiagnosis::diagnoseUnviableLookupResults(
     case MemberLookupResult::UR_LabelMismatch:
     case MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember:
     case MemberLookupResult::UR_ReferenceWritableKeyPathOnMutatingMember:
+    case MemberLookupResult::UR_KeyPathWithAnyObjectRootType:
       break;
     case MemberLookupResult::UR_UnavailableInExistential:
       diagnose(loc, diag::could_not_use_member_on_existential,

lib/Sema/CSDiagnostics.cpp

@@ -2399,6 +2399,24 @@ bool InaccessibleMemberFailure::diagnoseAsError() {
   return true;
 }
 
+bool AnyObjectKeyPathRootFailure::diagnoseAsError() {
+  // Diagnose use of AnyObject as root for a keypath
+
+  auto anchor = getAnchor();
+  auto loc = anchor->getLoc();
+  auto range = anchor->getSourceRange();
+
+  if (auto KPE = dyn_cast<KeyPathExpr>(anchor)) {
+    if (auto rootTyRepr = KPE->getRootType()) {
+      loc = rootTyRepr->getLoc();
+      range = rootTyRepr->getSourceRange();
+    }
+  }
+
+  emitDiagnostic(loc, diag::expr_swift_keypath_anyobject_root).highlight(range);
+  return true;
+}
+
 bool KeyPathSubscriptIndexHashableFailure::diagnoseAsError() {
   auto *anchor = getRawAnchor();
   auto *locator = getLocator();

lib/Sema/CSDiagnostics.h

@@ -987,6 +987,22 @@ class InaccessibleMemberFailure final : public FailureDiagnostic {
   bool diagnoseAsError() override;
 };
 
+
+// Diagnose an attempt to use AnyObject as the root type of a KeyPath
+//
+// ```swift
+// let keyPath = \AnyObject.bar
+// ```
+class AnyObjectKeyPathRootFailure final : public FailureDiagnostic {
+
+public:
+  AnyObjectKeyPathRootFailure(Expr *root, ConstraintSystem &cs,
+                              ConstraintLocator *locator)
+      : FailureDiagnostic(root, cs, locator) {}
+  
+  bool diagnoseAsError() override;
+};
+
 /// Diagnose an attempt to reference subscript as a keypath component
 /// where at least one of the index arguments doesn't conform to Hashable e.g.
 ///

lib/Sema/CSFix.cpp

@@ -397,6 +397,18 @@ AllowInaccessibleMember::create(ConstraintSystem &cs, ValueDecl *member,
   return new (cs.getAllocator()) AllowInaccessibleMember(cs, member, locator);
 }
 
+bool AllowAnyObjectKeyPathRoot::diagnose(Expr *root, bool asNote) const {
+  AnyObjectKeyPathRootFailure failure(root, getConstraintSystem(),
+                                      getLocator());
+  return failure.diagnose(asNote);
+}
+
+AllowAnyObjectKeyPathRoot *
+AllowAnyObjectKeyPathRoot::create(ConstraintSystem &cs,
+                                  ConstraintLocator *locator) {
+  return new (cs.getAllocator()) AllowAnyObjectKeyPathRoot(cs, locator);
+}
+
 bool TreatKeyPathSubscriptIndexAsHashable::diagnose(Expr *root,
                                                     bool asNote) const {
   KeyPathSubscriptIndexHashableFailure failure(root, getConstraintSystem(),

lib/Sema/CSFix.h

@@ -137,8 +137,11 @@ enum class FixKind : uint8_t {
   /// no access control.
   AllowInaccessibleMember,
 
+  /// Allow KeyPaths to use AnyObject as root type
+  AllowAnyObjectKeyPathRoot,
+  
   /// Using subscript references in the keypath requires that each
-  // of the index arguments to be Hashable.
+  /// of the index arguments to be Hashable.
   TreatKeyPathSubscriptIndexAsHashable,
 };
 
@@ -741,6 +744,22 @@ class AllowInaccessibleMember final : public ConstraintFix {
                                          ConstraintLocator *locator);
 };
 
+class AllowAnyObjectKeyPathRoot final : public ConstraintFix {
+
+  AllowAnyObjectKeyPathRoot(ConstraintSystem &cs, ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::AllowAnyObjectKeyPathRoot, locator) {}
+
+public:
+  std::string getName() const override {
+    return "allow anyobject as root type for a keypath";
+  }
+
+  bool diagnose(Expr *root, bool asNote = false) const override;
+
+  static AllowAnyObjectKeyPathRoot *create(ConstraintSystem &cs,
+                                           ConstraintLocator *locator);
+};
+
 class TreatKeyPathSubscriptIndexAsHashable final : public ConstraintFix {
   Type NonConformingType;
 

lib/Sema/CSGen.cpp

@@ -2907,8 +2907,10 @@ namespace {
 
       // For native key paths, traverse the key path components to set up
       // appropriate type relationships at each level.
+      auto rootLocator =
+          CS.getConstraintLocator(E, ConstraintLocator::KeyPathRoot);
       auto locator = CS.getConstraintLocator(E);
-      Type root = CS.createTypeVariable(locator);
+      Type root = CS.createTypeVariable(rootLocator);
 
       // If a root type was explicitly given, then resolve it now.
       if (auto rootRepr = E->getRootType()) {
@@ -3015,7 +3017,9 @@ namespace {
         base = optTy;
       }
 
-      auto rvalueBase = CS.createTypeVariable(locator);
+      auto baseLocator =
+          CS.getConstraintLocator(E, ConstraintLocator::KeyPathValue);
+      auto rvalueBase = CS.createTypeVariable(baseLocator);
       CS.addConstraint(ConstraintKind::Equal, base, rvalueBase, locator);
 
       // The result is a KeyPath from the root to the end component.

lib/Sema/CSSimplify.cpp

@@ -1851,6 +1851,22 @@ ConstraintSystem::matchTypesBindTypeVar(
     }
   }
 
+  // We do not allow keypaths to go through AnyObject. Let's create a fix
+  // so this can be diagnosed later.
+  if (auto loc = typeVar->getImpl().getLocator()) {
+    auto locPath = loc->getPath();
+
+    if (!locPath.empty() &&
+        locPath.back().getKind() == ConstraintLocator::KeyPathRoot &&
+        type->isAnyObject()) {
+      auto *fix = AllowAnyObjectKeyPathRoot::create(
+          *this, getConstraintLocator(locator));
+
+      if (recordFix(fix))
+        return getTypeMatchFailure(locator);
+    }
+  }
+
   // Okay. Bind below.
 
   // A constraint that binds any pointer to a void pointer is
@@ -3591,8 +3607,14 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
   if (memberName.isSimpleName() &&
       memberName.getBaseName().getKind() == DeclBaseName::Kind::Subscript &&
       !isKeyPathDynamicMemberLookup(memberLocator)) {
-    result.ViableCandidates.push_back(
-        OverloadChoice(baseTy, OverloadChoiceKind::KeyPathApplication));
+    if (baseTy->isAnyObject()) {
+      result.addUnviable(
+          OverloadChoice(baseTy, OverloadChoiceKind::KeyPathApplication),
+          MemberLookupResult::UR_KeyPathWithAnyObjectRootType);
+    } else {
+      result.ViableCandidates.push_back(
+          OverloadChoice(baseTy, OverloadChoiceKind::KeyPathApplication));
+    }
   }
 
   // If the base type is a tuple type, look for the named or indexed member
@@ -4273,6 +4295,8 @@ fixMemberRef(ConstraintSystem &cs, Type baseTy,
     case MemberLookupResult::UR_WritableKeyPathOnReadOnlyMember:
     case MemberLookupResult::UR_ReferenceWritableKeyPathOnMutatingMember:
       break;
+    case MemberLookupResult::UR_KeyPathWithAnyObjectRootType:
+      return AllowAnyObjectKeyPathRoot::create(cs, locator);
     }
   }
 
@@ -4950,7 +4974,7 @@ ConstraintSystem::simplifyKeyPathConstraint(Type keyPathTy,
         return SolutionKind::Error;
     }
   }
-  
+
   // See if we resolved overloads for all the components involved.
   enum {
     ReadOnly,
@@ -5085,7 +5109,7 @@ ConstraintSystem::simplifyKeyPathApplicationConstraint(
     }
     return SolutionKind::Unsolved;
   };
-  
+
   if (auto clas = keyPathTy->getAs<NominalType>()) {
     if (clas->getDecl() == getASTContext().getAnyKeyPathDecl()) {
       // Read-only keypath, whose projected value is upcast to `Any?`.
@@ -6288,6 +6312,7 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyFixConstraint(
   case FixKind::AllowClosureParameterDestructuring:
   case FixKind::MoveOutOfOrderArgument:
   case FixKind::AllowInaccessibleMember:
+  case FixKind::AllowAnyObjectKeyPathRoot:
   case FixKind::TreatKeyPathSubscriptIndexAsHashable:
     llvm_unreachable("handled elsewhere");
   }

lib/Sema/ConstraintLocator.cpp

@@ -79,6 +79,8 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
     case DynamicLookupResult:
     case ContextualType:
     case SynthesizedArgument:
+    case KeyPathRoot:
+    case KeyPathValue:
       if (unsigned numValues = numNumericValuesInPathElement(elt.getKind())) {
         id.AddInteger(elt.getValue());
         if (numValues > 1)
@@ -101,6 +103,26 @@ bool ConstraintLocator::isSubscriptMemberRef() const {
   return path.back().getKind() == ConstraintLocator::SubscriptMember;
 }
 
+bool ConstraintLocator::isKeyPathRoot() const {
+  auto *anchor = getAnchor();
+  auto path = getPath();
+
+  if (!anchor || path.empty())
+    return false;
+
+  return path.back().getKind() == ConstraintLocator::KeyPathRoot;
+}
+
+bool ConstraintLocator::isKeyPathValue() const {
+  auto *anchor = getAnchor();
+  auto path = getPath();
+
+  if (!anchor || path.empty())
+    return false;
+
+  return path.back().getKind() == ConstraintLocator::KeyPathValue;
+}
+
 bool ConstraintLocator::isResultOfKeyPathDynamicMemberLookup() const {
   return llvm::any_of(getPath(), [](const LocatorPathElt &elt) {
     return elt.isKeyPathDynamicMember();
@@ -310,6 +332,14 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) {
     case KeyPathDynamicMember:
       out << " keypath dynamic member lookup";
       break;
+
+    case KeyPathRoot:
+      out << " keypath root";
+      break;
+
+    case KeyPathValue:
+      out << " keypath value";
+      break;
     }
   }
 

lib/Sema/ConstraintLocator.h

@@ -129,6 +129,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     SynthesizedArgument,
     /// The member looked up via keypath based dynamic lookup.
     KeyPathDynamicMember,
+    /// The root of a keypath
+    KeyPathRoot,
+    /// The value of a keypath
+    KeyPathValue,
   };
 
   /// Determine the number of numeric values used for the given path
@@ -159,6 +163,8 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ImplicitlyUnwrappedDisjunctionChoice:
     case DynamicLookupResult:
     case ContextualType:
+    case KeyPathRoot:
+    case KeyPathValue:
       return 0;
 
     case OpenedGeneric:
@@ -225,6 +231,8 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case ContextualType:
     case SynthesizedArgument:
     case KeyPathDynamicMember:
+    case KeyPathRoot:
+    case KeyPathValue:
       return 0;
 
     case FunctionArgument:
@@ -522,6 +530,12 @@ class ConstraintLocator : public llvm::FoldingSetNode {
   /// e.g. `foo[0]` or `\Foo.[0]`
   bool isSubscriptMemberRef() const;
 
+  /// Determine whether given locator points to the keypath root
+  bool isKeyPathRoot() const;
+
+  /// 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;

lib/Sema/ConstraintSystem.h

@@ -855,25 +855,25 @@ struct MemberLookupResult {
   enum UnviableReason {
     /// Argument labels don't match.
     UR_LabelMismatch,
-    
+
     /// This uses a type like Self in its signature that cannot be used on an
     /// existential box.
     UR_UnavailableInExistential,
-    
+
     /// This is an instance member being accessed through something of metatype
     /// type.
     UR_InstanceMemberOnType,
-    
+
     /// This is a static/class member being accessed through an instance.
     UR_TypeMemberOnInstance,
-    
+
     /// This is a mutating member, being used on an rvalue.
     UR_MutatingMemberOnRValue,
-    
+
     /// The getter for this subscript or computed property is mutating and we
     /// only have an rvalue base.  This is more specific than the former one.
     UR_MutatingGetterOnRValue,
-    
+
     /// The member is inaccessible (e.g. a private member in another file).
     UR_Inaccessible,
 
@@ -882,11 +882,15 @@ struct MemberLookupResult {
     /// because it's not known upfront what access capability would the
     /// member have.
     UR_WritableKeyPathOnReadOnlyMember,
+
     /// This is a `ReferenceWritableKeyPath` being used to look up mutating
     /// member, used in situations involving dynamic member lookup via keypath,
     /// because it's not known upfront what access capability would the
     /// member have.
     UR_ReferenceWritableKeyPathOnMutatingMember,
+
+    /// This is a KeyPath whose root type is AnyObject
+    UR_KeyPathWithAnyObjectRootType
   };
 
   /// This is a list of considered (but rejected) candidates, along with a

test/expr/primary/keypath/keypath-objc.swift

@@ -143,3 +143,24 @@ func testParseErrors() {
 func testTypoCorrection() {
   let _: String = #keyPath(A.proString) // expected-error {{type 'A' has no member 'proString'}}
 }
+
+class SR_10146_1 {
+  @objc let b = 1
+}
+
+class SR_10146_2: SR_10146_1 {
+  let a = \AnyObject.b // expected-error {{the root type of a Swift key path cannot be 'AnyObject'}}
+}
+
+class SR_10146_3 {
+  @objc let abc: Int = 1
+  
+  func doNotCrash() {
+    let _: KeyPath<AnyObject, Int> = \.abc // expected-error {{the root type of a Swift key path cannot be 'AnyObject'}}
+  }
+
+  func doNotCrash_1(_ obj: AnyObject, _ kp: KeyPath<AnyObject, Int>) {
+    let _ = obj[keyPath: \.abc] // expected-error 2{{the root type of a Swift key path cannot be 'AnyObject'}}
+    let _ = obj[keyPath: kp] // expected-error {{the root type of a Swift key path cannot be 'AnyObject'}}
+  }
+}