[CS] Use custom locator element for callAsFunction

Introduce a `ImplicitCallAsFunction` locator path
element to represent an implicit member reference
to `callAsFunction`. Then adjust CSApply a little
to check whether it's finishing an apply for a
callable type, and if so build the implicit member
access.

lib/Sema/CSApply.cpp

@@ -6455,13 +6455,15 @@ static bool isValidDynamicCallableMethod(FuncDecl *method,
 // Build a reference to a `callAsFunction` method.
 static Expr *buildCallAsFunctionMethodRef(
     ExprRewriter &rewriter, ApplyExpr *apply, SelectedOverload selected,
-    ConstraintLocatorBuilder applyFunctionLoc) {
+    ConstraintLocator *calleeLoc) {
-  auto *fn = apply->getFn();
+  assert(calleeLoc->isLastElement<LocatorPathElt::ImplicitCallAsFunction>());
+  assert(cast<FuncDecl>(selected.choice.getDecl())->isCallAsFunctionMethod());
+
   // Create direct reference to `callAsFunction` method.
+  auto *fn = apply->getFn();
   auto *declRef = rewriter.buildMemberRef(
       fn, /*dotLoc*/ SourceLoc(), selected, DeclNameLoc(fn->getEndLoc()),
-      applyFunctionLoc, applyFunctionLoc, /*implicit*/ true,
+      calleeLoc, calleeLoc, /*implicit*/ true, AccessSemantics::Ordinary);
-      AccessSemantics::Ordinary);
   if (!declRef)
     return nullptr;
   declRef->setImplicit(apply->isImplicit());
@@ -6704,7 +6706,15 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
     callee = resolveConcreteDeclRef(decl, calleeLoc);
   }
 
-  // Resolve `callAsFunction` and `@dynamicCallable` applications.
+  // If this is an implicit call to a `callAsFunction` method, build the
+  // appropriate member reference.
+  if (cs.getType(fn)->getRValueType()->isCallableNominalType(dc)) {
+    fn = buildCallAsFunctionMethodRef(*this, apply, *overload, calleeLoc);
+    if (!fn)
+      return nullptr;
+  }
+
+  // Resolve a `@dynamicCallable` application.
   auto applyFunctionLoc =
       locator.withPathElement(ConstraintLocator::ApplyFunction);
   if (auto selected = solution.getOverloadChoiceIfAvailable(
@@ -6717,15 +6727,6 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
       if (isValidDynamicCallableMethod(method, methodType))
         return finishApplyDynamicCallable(
             apply, *selected, method, methodType, applyFunctionLoc);
-
-      // If this is an implicit call to a callAsFunction method, build the
-      // appropriate member reference.
-      if (method->isCallAsFunctionMethod()) {
-        fn = buildCallAsFunctionMethodRef(*this, apply, *selected,
-                                          applyFunctionLoc);
-        if (!fn)
-          return nullptr;
-      }
     }
   }
 

lib/Sema/CSSimplify.cpp

@@ -7391,18 +7391,16 @@ ConstraintSystem::simplifyApplicableFnConstraint(
   // is true.
   if (desugar2->isCallableNominalType(DC)) {
     auto memberLoc = getConstraintLocator(
-        outerLocator.withPathElement(ConstraintLocator::Member));
+        locator.withPathElement(ConstraintLocator::ImplicitCallAsFunction));
     // Add a `callAsFunction` member constraint, binding the member type to a
     // type variable.
     auto memberTy = createTypeVariable(memberLoc, /*options=*/0);
     // TODO: Revisit this if `static func callAsFunction` is to be supported.
     // Static member constraint requires `FunctionRefKind::DoubleApply`.
-    // TODO: Use a custom locator element to identify this member constraint
-    // instead of just pointing to the function expr.
     addValueMemberConstraint(origLValueType2,
                              DeclNameRef(ctx.Id_callAsFunction),
                              memberTy, DC, FunctionRefKind::SingleApply,
-                             /*outerAlternatives*/ {}, locator);
+                             /*outerAlternatives*/ {}, memberLoc);
     // Add new applicable function constraint based on the member type
     // variable.
     addConstraint(ConstraintKind::ApplicableFunction, func1, memberTy,

lib/Sema/ConstraintLocator.cpp

@@ -114,6 +114,7 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   case ConstraintLocator::KeyPathComponentResult:
   case ConstraintLocator::Condition:
   case ConstraintLocator::DynamicCallable:
+  case ConstraintLocator::ImplicitCallAsFunction:
     return 0;
 
   case ConstraintLocator::FunctionArgument:
@@ -458,6 +459,10 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) const {
     case DynamicCallable:
       out << "implicit call to @dynamicCallable method";
       break;
+
+    case ImplicitCallAsFunction:
+      out << "implicit reference to callAsFunction";
+      break;
     }
   }
   out << ']';

lib/Sema/ConstraintLocatorPathElts.def

@@ -76,6 +76,9 @@ SIMPLE_LOCATOR_PATH_ELT(FunctionResult)
 /// FIXME: Add support for named generic arguments?
 CUSTOM_LOCATOR_PATH_ELT(GenericArgument)
 
+/// An implicit reference to a 'callAsFunction' method of a nominal type.
+SIMPLE_LOCATOR_PATH_ELT(ImplicitCallAsFunction)
+
 /// Locator for a binding from an IUO disjunction choice.
 SIMPLE_LOCATOR_PATH_ELT(ImplicitlyUnwrappedDisjunctionChoice)
 

lib/Sema/ConstraintSystem.cpp

@@ -506,9 +506,11 @@ ConstraintSystem::getCalleeLocator(ConstraintLocator *locator,
     }
 
     // Handle an apply of a nominal type which supports callAsFunction.
-    if (fnTy->isCallableNominalType(DC))
+    if (fnTy->isCallableNominalType(DC)) {
-      return getConstraintLocator(anchor, ConstraintLocator::ApplyFunction);
+      return getConstraintLocator(anchor,
-
+                                  {LocatorPathElt::ApplyFunction(),
+                                   LocatorPathElt::ImplicitCallAsFunction()});
+    }
     return nullptr;
   };
 
@@ -3145,8 +3147,9 @@ void constraints::simplifyLocator(Expr *&anchor,
     case ConstraintLocator::LValueConversion:
     case ConstraintLocator::RValueAdjustment:
     case ConstraintLocator::UnresolvedMember:
-      // Arguments in autoclosure positions, lvalue and rvalue adjustments, and
+    case ConstraintLocator::ImplicitCallAsFunction:
-      // scalar-to-tuple conversions, and unresolved members are
+      // Arguments in autoclosure positions, lvalue and rvalue adjustments,
+      // unresolved members, and implicit callAsFunction references are
       // implicit.
       path = path.slice(1);
       continue;