[CS] Allow getCalleeLocator to find key path component callees

In order to do this we need it to take a ConstraintLocator argument so
we can tell which component we want the callee for. To make it clear
that we're looking for a callee at the anchor, also rename the member
to getAnchormostCalleeLocator.

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -4243,11 +4243,11 @@ namespace {
 
         auto locator = cs.getConstraintLocator(
             E, LocatorPathElt::KeyPathComponent(i));
-        if (kind == KeyPathExpr::Component::Kind::UnresolvedSubscript) {
-          locator =
-            cs.getConstraintLocator(locator,
-                                    ConstraintLocator::SubscriptMember);
-        }
+
+        // Adjust the locator such that it includes any additional elements to
+        // point to the component's callee, e.g a SubscriptMember for a
+        // subscript component.
+        locator = cs.getAnchormostCalleeLocator(locator);
 
         bool isDynamicMember = false;
         // If this is an unresolved link, make sure we resolved it.

lib/Sema/CSDiagnostics.cpp

@@ -123,7 +123,8 @@ Expr *FailureDiagnostic::getBaseExprFor(Expr *anchor) const {
 
 Optional<SelectedOverload> FailureDiagnostic::getChoiceFor(Expr *expr) const {
   auto &cs = getConstraintSystem();
-  return getOverloadChoiceIfAvailable(cs.getCalleeLocator(expr));
+  auto *loc = cs.getConstraintLocator(expr);
+  return getOverloadChoiceIfAvailable(cs.getAnchormostCalleeLocator(loc));
 }
 
 Type FailureDiagnostic::resolveInterfaceType(Type type,

lib/Sema/CSGen.cpp

@@ -3023,6 +3023,8 @@ namespace {
         }
 
         case KeyPathExpr::Component::Kind::TupleElement: {
+          // Note: If implemented, the logic in `getAnchormostCalleeLocator`
+          // will need updating to return the correct callee locator for this.
           llvm_unreachable("not implemented");
           break;
         }
@@ -3256,7 +3258,9 @@ namespace {
       // Record the labels.
       if (!labelsArePermanent)
         info.Labels = CS.allocateCopy(info.Labels);
-      CS.ArgumentInfos[CS.getArgumentInfoLocator(expr)] = info;
+
+      auto *locator = CS.getConstraintLocator(expr);
+      CS.ArgumentInfos[CS.getArgumentInfoLocator(locator)] = info;
     }
   };
 

lib/Sema/ConstraintSystem.cpp

@@ -413,14 +413,49 @@ ConstraintLocator *ConstraintSystem::getConstraintLocator(
   return getConstraintLocator(anchor, path, builder.getSummaryFlags());
 }
 
-ConstraintLocator *ConstraintSystem::getCalleeLocator(Expr *expr) {
+ConstraintLocator *
+ConstraintSystem::getAnchormostCalleeLocator(ConstraintLocator *locator) {
+  auto *anchor = locator->getAnchor();
+  assert(anchor && "Expected an anchor!");
+
+  // If we have a locator that starts with a key path component element, we
+  // may have a callee given by a property or subscript component.
+  if (auto componentElt =
+          locator->getFirstElementAs<LocatorPathElt::KeyPathComponent>()) {
+    auto *kpExpr = cast<KeyPathExpr>(anchor);
+    auto component = kpExpr->getComponents()[componentElt->getIndex()];
+
+    using ComponentKind = KeyPathExpr::Component::Kind;
+    switch (component.getKind()) {
+    case ComponentKind::UnresolvedSubscript:
+    case ComponentKind::Subscript:
+      // For a subscript the callee is given by 'component -> subscript member'.
+      return getConstraintLocator(
+          anchor, {*componentElt, ConstraintLocator::SubscriptMember});
+    case ComponentKind::UnresolvedProperty:
+    case ComponentKind::Property:
+      // For a property, the choice is just given by the component.
+      return getConstraintLocator(anchor, *componentElt);
+    case ComponentKind::TupleElement:
+      llvm_unreachable("Not implemented by CSGen");
+      break;
+    case ComponentKind::Invalid:
+    case ComponentKind::OptionalForce:
+    case ComponentKind::OptionalChain:
+    case ComponentKind::OptionalWrap:
+    case ComponentKind::Identity:
+      // These components don't have any callee associated, so just continue.
+      break;
+    }
+  }
+
   // Make sure we handle subscripts before looking at apply exprs. We don't
   // want to return a subscript member locator for an expression such as x[](y),
   // as its callee is not the subscript, but rather the function it returns.
-  if (isa<SubscriptExpr>(expr))
-    return getConstraintLocator(expr, ConstraintLocator::SubscriptMember);
+  if (isa<SubscriptExpr>(anchor))
+    return getConstraintLocator(anchor, ConstraintLocator::SubscriptMember);
 
-  if (auto *applyExpr = dyn_cast<ApplyExpr>(expr)) {
+  if (auto *applyExpr = dyn_cast<ApplyExpr>(anchor)) {
     auto *fnExpr = applyExpr->getFn();
     // For an apply of a metatype, we have a short-form constructor. Unlike
     // other locators to callees, these are anchored on the apply expression
@@ -436,27 +471,27 @@ ConstraintLocator *ConstraintSystem::getCalleeLocator(Expr *expr) {
     // Otherwise fall through and look for locators anchored on the function
     // expr. For CallExprs, this can look through things like parens and
     // optional chaining.
-    if (auto *callExpr = dyn_cast<CallExpr>(expr)) {
-      expr = callExpr->getDirectCallee();
+    if (auto *callExpr = dyn_cast<CallExpr>(anchor)) {
+      anchor = callExpr->getDirectCallee();
     } else {
-      expr = fnExpr;
+      anchor = fnExpr;
     }
   }
 
-  if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {
+  if (auto *UDE = dyn_cast<UnresolvedDotExpr>(anchor)) {
     return getConstraintLocator(
-        expr, TC.getSelfForInitDelegationInConstructor(DC, UDE)
+        anchor, TC.getSelfForInitDelegationInConstructor(DC, UDE)
                   ? ConstraintLocator::ConstructorMember
                   : ConstraintLocator::Member);
   }
 
-  if (isa<UnresolvedMemberExpr>(expr))
-    return getConstraintLocator(expr, ConstraintLocator::UnresolvedMember);
+  if (isa<UnresolvedMemberExpr>(anchor))
+    return getConstraintLocator(anchor, ConstraintLocator::UnresolvedMember);
 
-  if (isa<MemberRefExpr>(expr))
-    return getConstraintLocator(expr, ConstraintLocator::Member);
+  if (isa<MemberRefExpr>(anchor))
+    return getConstraintLocator(anchor, ConstraintLocator::Member);
 
-  return getConstraintLocator(expr);
+  return getConstraintLocator(anchor);
 }
 
 Type ConstraintSystem::openUnboundGenericType(UnboundGenericType *unbound,
@@ -2462,7 +2497,7 @@ bool ConstraintSystem::diagnoseAmbiguityWithFixes(
       return false;
 
     const auto *fix = fixes.front();
-    auto *calleeLocator = getCalleeLocator(fix->getAnchor());
+    auto *calleeLocator = getAnchormostCalleeLocator(fix->getLocator());
     if (commonCalleeLocator && commonCalleeLocator != calleeLocator)
       return false;
 
@@ -2829,7 +2864,9 @@ void ConstraintSystem::generateConstraints(
   }
 }
 
-ConstraintLocator *ConstraintSystem::getArgumentInfoLocator(Expr *anchor) {
+ConstraintLocator *
+ConstraintSystem::getArgumentInfoLocator(ConstraintLocator *locator) {
+  auto *anchor = locator->getAnchor();
   if (!anchor)
     return nullptr;
 
@@ -2838,12 +2875,12 @@ ConstraintLocator *ConstraintSystem::getArgumentInfoLocator(Expr *anchor) {
     return getConstraintLocator(fnExpr);
   }
 
-  return getCalleeLocator(anchor);
+  return getAnchormostCalleeLocator(locator);
 }
 
 Optional<ConstraintSystem::ArgumentInfo>
 ConstraintSystem::getArgumentInfo(ConstraintLocator *locator) {
-  if (auto *infoLocator = getArgumentInfoLocator(locator->getAnchor())) {
+  if (auto *infoLocator = getArgumentInfoLocator(locator)) {
     auto known = ArgumentInfos.find(infoLocator);
     if (known != ArgumentInfos.end())
       return known->second;

lib/Sema/ConstraintSystem.h

@@ -1579,9 +1579,10 @@ class ConstraintSystem {
   /// that locator.
   llvm::DenseMap<ConstraintLocator *, ArgumentInfo> ArgumentInfos;
 
-  /// Form a locator with given anchor which then could be used
-  /// to retrieve argument information cached in the constraint system.
-  ConstraintLocator *getArgumentInfoLocator(Expr *anchor);
+  /// Form a locator that can be used to retrieve argument information cached in
+  /// the constraint system for the callee described by the anchor of the
+  /// passed locator.
+  ConstraintLocator *getArgumentInfoLocator(ConstraintLocator *locator);
 
   /// Retrieve the argument info that is associated with a member
   /// reference at the given locator.
@@ -2008,10 +2009,21 @@ class ConstraintSystem {
     return e != ExprWeights.end() ? e->second.second : nullptr;
   }
 
-  /// Returns a locator describing the callee for a given expression. For
-  /// a function application, this is a locator describing the function expr.
-  /// For an unresolved dot/member, this is a locator to the member.
-  ConstraintLocator *getCalleeLocator(Expr *expr);
+  /// Returns a locator describing the callee for the anchor of a given locator.
+  ///
+  /// - For an unresolved dot/member anchor, this will be a locator describing
+  /// the member.
+  ///
+  /// - For a subscript anchor, this will be a locator describing the subscript
+  /// member.
+  ///
+  /// - For a key path anchor with a property/subscript component path element,
+  /// this will be a locator describing the decl referenced by the component.
+  ///
+  /// - For a function application anchor, this will be a locator describing the
+  /// 'direct callee' of the call. For example, for the expression \c x.foo?()
+  /// the returned locator will describe the member \c foo.
+  ConstraintLocator *getAnchormostCalleeLocator(ConstraintLocator *locator);
 
 public: