[CS] Fix resolution of key path component callees

Previously we were just searching for a resolved
overload with the same component index and anchor
for the key path. However unfortunately that found
dynamic member lookup overloads in certain cases,
leading to an incorrect mutability capability.

Change the logic such that it uses the callee
locator for a given key path component, ensuring
we only match against the "direct" callee, and
not any dynamic members it might be also referring
to.

Resolves SR-11896.

Author: hamishknight

lib/Sema/CSSimplify.cpp

@@ -6729,22 +6729,6 @@ ConstraintSystem::simplifyKeyPathConstraint(
   // The constraint ought to have been anchored on a KeyPathExpr.
   auto keyPath = cast<KeyPathExpr>(locator.getBaseLocator()->getAnchor());
 
-  // Gather overload choices for any key path components associated with this
-  // key path.
-  SmallVector<OverloadChoice, 4> choices;
-  choices.resize(keyPath->getComponents().size());
-  for (auto resolvedItem = resolvedOverloadSets; resolvedItem;
-       resolvedItem = resolvedItem->Previous) {
-    auto locator = resolvedItem->Locator;
-    auto path = locator->getPath();
-    if (locator->getAnchor() != keyPath || path.size() > 2)
-      continue;
-
-    if (auto kpElt = path[0].getAs<LocatorPathElt::KeyPathComponent>()) {
-      choices[kpElt->getIndex()] = resolvedItem->Choice;
-    }
-  }
-
   keyPathTy = getFixedTypeRecursive(keyPathTy, /*want rvalue*/ true);
   bool definitelyFunctionType = false;
   bool definitelyKeyPathType = false;
@@ -6806,6 +6790,30 @@ ConstraintSystem::simplifyKeyPathConstraint(
       return SolutionKind::Error;
   }
 
+  // First gather the callee locators for the individual components.
+  SmallVector<std::pair<ConstraintLocator *, OverloadChoice>, 4> choices;
+  for (unsigned i : indices(keyPath->getComponents())) {
+    auto *componentLoc = getConstraintLocator(
+        locator.withPathElement(LocatorPathElt::KeyPathComponent(i)));
+    auto *calleeLoc = getCalleeLocator(componentLoc);
+    choices.push_back({calleeLoc, OverloadChoice()});
+  }
+
+  // Then search for the resolved overloads.
+  for (auto resolvedItem = resolvedOverloadSets; resolvedItem;
+       resolvedItem = resolvedItem->Previous) {
+    auto locator = resolvedItem->Locator;
+    auto kpElt = locator->getFirstElementAs<LocatorPathElt::KeyPathComponent>();
+    if (!kpElt || locator->getAnchor() != keyPath)
+      continue;
+
+    auto &choice = choices[kpElt->getIndex()];
+    if (choice.first == locator) {
+      assert(choice.second.isInvalid() && "Resolved same component twice?");
+      choice.second = resolvedItem->Choice;
+    }
+  }
+
   // See if we resolved overloads for all the components involved.
   enum {
     ReadOnly,
@@ -6827,12 +6835,15 @@ ConstraintSystem::simplifyKeyPathConstraint(
     case KeyPathExpr::Component::Kind::Subscript:
     case KeyPathExpr::Component::Kind::UnresolvedProperty:
     case KeyPathExpr::Component::Kind::UnresolvedSubscript: {
+      auto *calleeLoc = choices[i].first;
+      auto choice = choices[i].second;
+
       // If no choice was made, leave the constraint unsolved. But when
       // generating constraints, we may already have enough information
       // to determine whether the result will be a function type vs BGT KeyPath
       // type, so continue through components to create new constraint at the
       // end.
-      if (choices[i].isInvalid() || anyComponentsUnresolved) {
+      if (choice.isInvalid() || anyComponentsUnresolved) {
         if (flags.contains(TMF_GenerateConstraints)) {
           anyComponentsUnresolved = true;
           continue;
@@ -6864,23 +6875,20 @@ ConstraintSystem::simplifyKeyPathConstraint(
       }
 
       // tuple elements do not change the capability of the key path
-      if (choices[i].getKind() == OverloadChoiceKind::TupleIndex) {
+      if (choice.getKind() == OverloadChoiceKind::TupleIndex) {
         continue;
       }
 
       // Discarded unsupported non-decl member lookups.
-      if (!choices[i].isDecl()) {
+      if (!choice.isDecl()) {
         return SolutionKind::Error;
       }
 
-      auto storage = dyn_cast<AbstractStorageDecl>(choices[i].getDecl());
-
-      auto *componentLoc = getConstraintLocator(
-          locator.withPathElement(LocatorPathElt::KeyPathComponent(i)));
+      auto storage = dyn_cast<AbstractStorageDecl>(choice.getDecl());
 
       if (auto *fix = AllowInvalidRefInKeyPath::forRef(
-              *this, choices[i].getDecl(), componentLoc)) {
-        if (!hasFixFor(componentLoc, FixKind::AllowTypeOrInstanceMember))
+              *this, choice.getDecl(), calleeLoc)) {
+        if (!hasFixFor(calleeLoc, FixKind::AllowTypeOrInstanceMember))
           if (!shouldAttemptFixes() || recordFix(fix))
             return SolutionKind::Error;
 

test/Constraints/keypath_dynamic_member_lookup.swift

@@ -410,12 +410,51 @@ func testIUOUnwrap(_ x: SR_11893) {
   // CHECK: keypath $KeyPath<SR_11893_Base, Int>, (root $SR_11893_Base; gettable_property $Optional<Int>,  id @$s29keypath_dynamic_member_lookup13SR_11893_BaseVySiSgSicig : $@convention(method) (Int, SR_11893_Base) -> Optional<Int>, getter @$s29keypath_dynamic_member_lookup13SR_11893_BaseVySiSgSicipACTK : $@convention(thin) (@in_guaranteed SR_11893_Base, UnsafeRawPointer) -> @out Optional<Int>, indices [%$0 : $Int : $Int], indices_equals @$sSiTH : $@convention(thin) (UnsafeRawPointer, UnsafeRawPointer) -> Bool, indices_hash @$sSiTh : $@convention(thin) (UnsafeRawPointer) -> Int; optional_force : $Int)
   x[5]
 
-  // FIXME(SR-11896): We should not vend a WritableKeyPath for the "outer" key path here.
   // CHECK: [[INNER_KP:%[0-9]+]] = keypath $KeyPath<SR_11893_Base, Int>, (root $SR_11893_Base; stored_property #SR_11893_Base.i : $Optional<Int>; optional_force : $Int)
-  // CHECK: keypath $WritableKeyPath<SR_11893, ()>, (root $SR_11893; gettable_property $(),  id @$s29keypath_dynamic_member_lookup8SR_11893V0B6Memberys7KeyPathCyAA0E11_11893_BaseVSiG_tcig : $@convention(method) (@guaranteed KeyPath<SR_11893_Base, Int>, SR_11893) -> (), getter @$s29keypath_dynamic_member_lookup8SR_11893V0B6Memberys7KeyPathCyAA0E11_11893_BaseVSiG_tcipACTK : $@convention(thin) (@in_guaranteed SR_11893, UnsafeRawPointer) -> @out (), indices [%$0 : $KeyPath<SR_11893_Base, Int> : $KeyPath<SR_11893_Base, Int>], indices_equals @$ss7KeyPathCy29keypath_dynamic_member_lookup13SR_11893_BaseVSiGTH : $@convention(thin) (UnsafeRawPointer, UnsafeRawPointer) -> Bool, indices_hash @$ss7KeyPathCy29keypath_dynamic_member_lookup13SR_11893_BaseVSiGTh : $@convention(thin) (UnsafeRawPointer) -> Int) ([[INNER_KP]])
+  // CHECK: keypath $KeyPath<SR_11893, ()>, (root $SR_11893; gettable_property $(),  id @$s29keypath_dynamic_member_lookup8SR_11893V0B6Memberys7KeyPathCyAA0E11_11893_BaseVSiG_tcig : $@convention(method) (@guaranteed KeyPath<SR_11893_Base, Int>, SR_11893) -> (), getter @$s29keypath_dynamic_member_lookup8SR_11893V0B6Memberys7KeyPathCyAA0E11_11893_BaseVSiG_tcipACTK : $@convention(thin) (@in_guaranteed SR_11893, UnsafeRawPointer) -> @out (), indices [%$0 : $KeyPath<SR_11893_Base, Int> : $KeyPath<SR_11893_Base, Int>], indices_equals @$ss7KeyPathCy29keypath_dynamic_member_lookup13SR_11893_BaseVSiGTH : $@convention(thin) (UnsafeRawPointer, UnsafeRawPointer) -> Bool, indices_hash @$ss7KeyPathCy29keypath_dynamic_member_lookup13SR_11893_BaseVSiGTh : $@convention(thin) (UnsafeRawPointer) -> Int) ([[INNER_KP]])
   _ = \SR_11893.i
 
   // CHECK: [[INNER_SUB_KP:%[0-9]+]] = keypath $KeyPath<SR_11893_Base, Int>, (root $SR_11893_Base; gettable_property $Optional<Int>,  id @$s29keypath_dynamic_member_lookup13SR_11893_BaseVySiSgSicig : $@convention(method) (Int, SR_11893_Base) -> Optional<Int>, getter @$s29keypath_dynamic_member_lookup13SR_11893_BaseVySiSgSicipACTK : $@convention(thin) (@in_guaranteed SR_11893_Base, UnsafeRawPointer) -> @out Optional<Int>, indices [%$0 : $Int : $Int], indices_equals @$sSiTH : $@convention(thin) (UnsafeRawPointer, UnsafeRawPointer) -> Bool, indices_hash @$sSiTh : $@convention(thin) (UnsafeRawPointer) -> Int; optional_force : $Int)
   // CHECK: keypath $KeyPath<SR_11893, ()>, (root $SR_11893; gettable_property $(),  id @$s29keypath_dynamic_member_lookup8SR_11893V0B6Memberys7KeyPathCyAA0E11_11893_BaseVSiG_tcig : $@convention(method) (@guaranteed KeyPath<SR_11893_Base, Int>, SR_11893) -> (), getter @$s29keypath_dynamic_member_lookup8SR_11893V0B6Memberys7KeyPathCyAA0E11_11893_BaseVSiG_tcipACTK : $@convention(thin) (@in_guaranteed SR_11893, UnsafeRawPointer) -> @out (), indices [%$0 : $KeyPath<SR_11893_Base, Int> : $KeyPath<SR_11893_Base, Int>], indices_equals @$ss7KeyPathCy29keypath_dynamic_member_lookup13SR_11893_BaseVSiGTH : $@convention(thin) (UnsafeRawPointer, UnsafeRawPointer) -> Bool, indices_hash @$ss7KeyPathCy29keypath_dynamic_member_lookup13SR_11893_BaseVSiGTh : $@convention(thin) (UnsafeRawPointer) -> Int) ([[INNER_SUB_KP]])
   _ = \SR_11893.[5]
 }
+
+// SR-11896: Make sure the outer key path reflects the mutability of the 'dynamicMember:' subscript.
+struct SR_11896_Base {
+  var mutable: Int
+  let immutable: Int
+}
+
+@dynamicMemberLookup
+struct SR_11896_Mutable {
+  subscript(dynamicMember kp: KeyPath<SR_11896_Base, Int>) -> Int {
+    get { 5 } set {}
+  }
+}
+
+@dynamicMemberLookup
+struct SR_11896_Immutable {
+  subscript(dynamicMember kp: KeyPath<SR_11896_Base, Int>) -> Int {
+    get { 5 }
+  }
+}
+
+
+// CHECK-LABEL: sil hidden @$s29keypath_dynamic_member_lookup21testKeyPathMutabilityyyF : $@convention(thin) () -> ()
+func testKeyPathMutability() {
+  // CHECK: keypath $KeyPath<SR_11896_Base, Int>, (root $SR_11896_Base; stored_property #SR_11896_Base.mutable : $Int)
+  // CHECK: keypath $WritableKeyPath<SR_11896_Mutable, Int>, (root $SR_11896_Mutable; settable_property $Int
+  _ = \SR_11896_Mutable.mutable
+
+  // CHECK: keypath $KeyPath<SR_11896_Base, Int>, (root $SR_11896_Base; stored_property #SR_11896_Base.immutable : $Int)
+  // CHECK: keypath $WritableKeyPath<SR_11896_Mutable, Int>, (root $SR_11896_Mutable; settable_property $Int
+  _ = \SR_11896_Mutable.immutable
+
+  // CHECK: keypath $KeyPath<SR_11896_Base, Int>, (root $SR_11896_Base; stored_property #SR_11896_Base.mutable : $Int)
+  // CHECK: keypath $KeyPath<SR_11896_Immutable, Int>, (root $SR_11896_Immutable; gettable_property $Int
+  _ = \SR_11896_Immutable.mutable
+
+  // CHECK: keypath $KeyPath<SR_11896_Base, Int>, (root $SR_11896_Base; stored_property #SR_11896_Base.immutable : $Int)
+  // CHECK: keypath $KeyPath<SR_11896_Immutable, Int>, (root $SR_11896_Immutable; gettable_property $Int
+  _ = \SR_11896_Immutable.immutable
+}