[CS] Resolve callees for callAsFunction .member exprs (#28604)

[CS] Resolve callees for callAsFunction `.member` exprs

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -2385,8 +2385,12 @@ namespace {
 
       // If there was an argument, apply it.
       if (auto arg = expr->getArgument()) {
-        auto callee = resolveConcreteDeclRef(selected.choice.getDecl(),
-                                             memberLocator);
+        // Find the callee. Note this may be different to the member being
+        // referenced for things like callAsFunction.
+        auto *calleeLoc = cs.getCalleeLocator(exprLoc);
+        auto calleeOverload = solution.getOverloadChoice(calleeLoc);
+        auto callee = resolveConcreteDeclRef(calleeOverload.choice.getDecl(),
+                                             calleeLoc);
         ApplyExpr *apply = CallExpr::create(
             ctx, result, arg, expr->getArgumentLabels(),
             expr->getArgumentLabelLocs(), expr->hasTrailingClosure(),

lib/Sema/ConstraintSystem.cpp

@@ -475,31 +475,38 @@ ConstraintSystem::getCalleeLocator(ConstraintLocator *locator,
   if (isa<SubscriptExpr>(anchor))
     return getConstraintLocator(anchor, ConstraintLocator::SubscriptMember);
 
+  auto getSpecialFnCalleeLoc = [&](Type fnTy) -> ConstraintLocator * {
+    // FIXME: We should probably assert that we don't get a type variable
+    // here to make sure we only retrieve callee locators for resolved calls,
+    // ensuring that callee locators don't change after binding a type.
+    // Unfortunately CSDiag currently calls into getCalleeLocator, so all bets
+    // are off. Once we remove that legacy diagnostic logic, we should be able
+    // to assert here.
+    fnTy = getFixedTypeRecursive(fnTy, /*wantRValue*/ true);
+
+    // For an apply of a metatype, we have a short-form constructor. Unlike
+    // other locators to callees, these are anchored on the apply expression
+    // rather than the function expr.
+    if (fnTy->is<AnyMetatypeType>()) {
+      return getConstraintLocator(anchor,
+                                  {LocatorPathElt::ApplyFunction(),
+                                   LocatorPathElt::ConstructorMember()});
+    }
+
+    // Handle an apply of a nominal type which supports callAsFunction.
+    if (fnTy->isCallableNominalType(DC))
+      return getConstraintLocator(anchor, ConstraintLocator::ApplyFunction);
+
+    return nullptr;
+  };
+
   if (lookThroughApply) {
     if (auto *applyExpr = dyn_cast<ApplyExpr>(anchor)) {
       auto *fnExpr = applyExpr->getFn();
 
-      // FIXME: We should probably assert that we don't get a type variable
-      // here to make sure we only retrieve callee locators for resolved calls,
-      // ensuring that callee locators don't change after binding a type.
-      // Unfortunately CSDiag currently calls into getCalleeLocator, so all bets
-      // are off. Once we remove that legacy diagnostic logic, we should be able
-      // to assert here.
-      auto fnTy = getFixedTypeRecursive(getType(fnExpr), /*wantRValue*/ true);
-
-      // For an apply of a metatype, we have a short-form constructor. Unlike
-      // other locators to callees, these are anchored on the apply expression
-      // rather than the function expr.
-      if (fnTy->is<AnyMetatypeType>()) {
-        auto *fnLocator =
-            getConstraintLocator(applyExpr, ConstraintLocator::ApplyFunction);
-        return getConstraintLocator(fnLocator,
-                                    ConstraintLocator::ConstructorMember);
-      }
-
-      // Handle an apply of a nominal type which supports callAsFunction.
-      if (fnTy->isCallableNominalType(DC))
-        return getConstraintLocator(anchor, ConstraintLocator::ApplyFunction);
+      // Handle special cases for applies of non-function types.
+      if (auto *loc = getSpecialFnCalleeLoc(getType(fnExpr)))
+        return loc;
 
       // Otherwise fall through and look for locators anchored on the function
       // expr. For CallExprs, this can look through things like parens and
@@ -519,8 +526,23 @@ ConstraintSystem::getCalleeLocator(ConstraintLocator *locator,
                     : ConstraintLocator::Member);
   }
 
-  if (isa<UnresolvedMemberExpr>(anchor))
-    return getConstraintLocator(anchor, ConstraintLocator::UnresolvedMember);
+  if (auto *UME = dyn_cast<UnresolvedMemberExpr>(anchor)) {
+    auto *calleeLoc =
+        getConstraintLocator(UME, ConstraintLocator::UnresolvedMember);
+
+    // Handle special cases for applies of non-function types.
+    // FIXME: Consider re-designing the AST such that an unresolved member expr
+    // with arguments uses a CallExpr, which would make this logic unnecessary
+    // and clean up a bunch of other special cases. Doing so may require a bit
+    // of hacking in CSGen though.
+    if (UME->hasArguments()) {
+      if (auto overload = findSelectedOverloadFor(calleeLoc)) {
+        if (auto *loc = getSpecialFnCalleeLoc(overload->boundType))
+          return loc;
+      }
+    }
+    return calleeLoc;
+  }
 
   if (isa<MemberRefExpr>(anchor))
     return getConstraintLocator(anchor, ConstraintLocator::Member);
@@ -3356,11 +3378,17 @@ ConstraintSystem::getArgumentInfoLocator(ConstraintLocator *locator) {
   if (!anchor)
     return nullptr;
 
+  // Applies and unresolved member exprs can have callee locators that are
+  // dependent on the type of their function, which may not have been resolved
+  // yet. Therefore we need to handle them specially.
   if (auto *apply = dyn_cast<ApplyExpr>(anchor)) {
     auto *fnExpr = getArgumentLabelTargetExpr(apply->getFn());
     return getConstraintLocator(fnExpr);
   }
 
+  if (auto *UME = dyn_cast<UnresolvedMemberExpr>(anchor))
+    return getConstraintLocator(UME);
+
   return getCalleeLocator(locator);
 }
 

test/SILGen/default_arguments.swift

@@ -431,3 +431,13 @@ func testCallableWithDefault(_ x: CallableWithDefault) {
   // CHECK: apply [[CALL_AS_FN]]([[I]], [[STR]], {{%[0-9]+}})
   x(y: 5)
 }
+
+// FIXME: Arguably we shouldn't allow calling a constructor like this, as
+// we usually require the user write an explicit '.init'.
+struct WeirdUMEInitCase {
+  static let ty = WeirdUMEInitCase.self
+  init(_ x: Int = 0) {}
+}
+
+let _: WeirdUMEInitCase = .ty()
+let _: WeirdUMEInitCase = .ty(5)

test/Sema/call_as_function_simple.swift

@@ -228,3 +228,13 @@ struct PrivateCallable {
 func testAccessControl(_ x: PrivateCallable) {
   x(5) // expected-error {{'callAsFunction' is inaccessible due to 'private' protection level}}
 }
+
+struct SR_11909 {
+  static let s = SR_11909()
+  func callAsFunction(_ x: Int = 0) -> SR_11909 { SR_11909() }
+}
+
+func testDefaultsWithUMEs(_ x: SR_11909) {
+  let _: SR_11909 = .s()
+  let _: SR_11909 = .s(5)
+}