[CSDiagnostics] Use getAnchormostCalleeLocator

Have FailureDiagnostic::getChoiceFor take a ConstraintLocator argument
which is passed through to getAnchormostCalleeLocator, and rename to
getAnchormostChoiceFor to make the semantics clear. In addition, add
a convenience getAnchormostChoice member for the common case of getting
the choice for the anchor of the failure's locator.

This change means we can now resolve callees for failures associated
with key path subscript components.

Resolves SR-11435.

Author: hamishknight

lib/Sema/CSDiagnostics.cpp

@@ -121,10 +121,14 @@ Expr *FailureDiagnostic::getBaseExprFor(Expr *anchor) const {
   return nullptr;
 }
 
-Optional<SelectedOverload> FailureDiagnostic::getChoiceFor(Expr *expr) const {
+Optional<SelectedOverload> FailureDiagnostic::getAnchormostChoice() const {
+  return getAnchormostChoiceFor(getLocator());
+}
+
+Optional<SelectedOverload>
+FailureDiagnostic::getAnchormostChoiceFor(ConstraintLocator *locator) const {
   auto &cs = getConstraintSystem();
-  auto *loc = cs.getConstraintLocator(expr);
-  return getOverloadChoiceIfAvailable(cs.getAnchormostCalleeLocator(loc));
+  return getOverloadChoiceIfAvailable(cs.getAnchormostCalleeLocator(locator));
 }
 
 Type FailureDiagnostic::resolveInterfaceType(Type type,
@@ -159,7 +163,8 @@ Type FailureDiagnostic::resolveInterfaceType(Type type,
 }
 
 /// Given an apply expr, returns true if it is expected to have a direct callee
-/// overload, resolvable using `getChoiceFor`. Otherwise, returns false.
+/// overload, resolvable using `getAnchormostChoiceFor`. Otherwise, returns
+/// false.
 static bool shouldHaveDirectCalleeOverload(const CallExpr *callExpr) {
   auto *fnExpr = callExpr->getDirectCallee();
 
@@ -221,7 +226,7 @@ FailureDiagnostic::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
 
   ValueDecl *callee = nullptr;
   Type rawFnType;
-  if (auto overload = getChoiceFor(anchor)) {
+  if (auto overload = getAnchormostChoiceFor(argLocator)) {
     // If we have resolved an overload for the callee, then use that to get the
     // function type and callee.
     callee = overload->choice.getDeclOrNull();
@@ -375,7 +380,7 @@ ValueDecl *RequirementFailure::getDeclRef() const {
   if (isFromContextualType())
     return getAffectedDeclFromType(cs.getContextualType());
 
-  if (auto overload = getChoiceFor(getRawAnchor())) {
+  if (auto overload = getAnchormostChoice()) {
     // If there is a declaration associated with this
     // failure e.g. an overload choice of the call
     // expression, let's see whether failure is
@@ -747,7 +752,7 @@ bool LabelingFailure::diagnoseAsNote() {
     return "(" + str + ")";
   };
 
-  auto selectedOverload = getChoiceFor(anchor);
+  auto selectedOverload = getAnchormostChoice();
   if (!selectedOverload)
     return false;
 
@@ -3121,8 +3126,8 @@ bool AllowTypeOrInstanceMemberFailure::diagnoseAsError() {
         return true;
       };
 
-      auto selection = getChoiceFor(ctorRef->getBase());
-      if (selection) {
+      auto *baseLoc = cs.getConstraintLocator(ctorRef->getBase());
+      if (auto selection = getAnchormostChoiceFor(baseLoc)) {
         OverloadChoice choice = selection->choice;
         if (choice.isDecl() && isMutable(choice.getDecl()) &&
             !isCallArgument(initCall) &&
@@ -4282,15 +4287,13 @@ bool MutatingMemberRefOnImmutableBase::diagnoseAsError() {
 }
 
 bool InvalidTupleSplatWithSingleParameterFailure::diagnoseAsError() {
-  auto *anchor = getRawAnchor();
-
-  auto selectedOverload = getChoiceFor(anchor);
+  auto selectedOverload = getAnchormostChoice();
   if (!selectedOverload || !selectedOverload->choice.isDecl())
     return false;
 
   auto *choice = selectedOverload->choice.getDecl();
 
-  auto *argExpr = getArgumentExprFor(anchor);
+  auto *argExpr = getArgumentExprFor(getRawAnchor());
   if (!argExpr)
     return false;
 

lib/Sema/CSDiagnostics.h

@@ -176,7 +176,13 @@ class FailureDiagnostic {
   /// reference or subscript, nullptr otherwise.
   Expr *getArgumentExprFor(Expr *anchor) const;
 
-  Optional<SelectedOverload> getChoiceFor(Expr *) const;
+  /// \returns The overload choice made by the constraint system for the callee
+  /// of this failure's raw anchor, or \c None if no such choice can be found.
+  Optional<SelectedOverload> getAnchormostChoice() const;
+
+  /// \returns The overload choice made by the constraint system for the callee
+  /// of a given locator's anchor, or \c None if no such choice can be found.
+  Optional<SelectedOverload> getAnchormostChoiceFor(ConstraintLocator *) const;
 
   /// For a given locator describing a function argument conversion, or a
   /// constraint within an argument conversion, returns information about the

test/Constraints/generics.swift

@@ -819,3 +819,16 @@ func test_correct_identification_of_requirement_source() {
   _ = X(A(), 17)
   // expected-error@-1 {{initializer 'init(_:_:)' requires that 'Int' conform to 'P'}}
 }
+
+struct SR11435<T> {
+  subscript<U : P & Hashable>(x x: U) -> U { x } // expected-note {{where 'U' = 'Int'}}
+}
+
+extension SR11435 where T : P { // expected-note {{where 'T' = 'Int'}}
+  var foo: Int { 0 }
+}
+
+func test_identification_of_key_path_component_callees() {
+  _ = \SR11435<Int>.foo // expected-error {{property 'foo' requires that 'Int' conform to 'P'}}
+  _ = \SR11435<Int>.[x: 5] // expected-error {{subscript 'subscript(x:)' requires that 'Int' conform to 'P'}}
+}