[CS] Allow getCalleeLocator to look through optional chaining

Author: hamishknight

lib/Sema/CSDiagnostics.cpp

@@ -224,20 +224,23 @@ FailureDiagnostic::getFunctionArgApplyInfo(ConstraintLocator *locator) const {
   if (auto overload = getChoiceFor(anchor)) {
     // If we have resolved an overload for the callee, then use that to get the
     // function type and callee.
-    if (auto *decl = overload->choice.getDeclOrNull())
-      callee = decl;
-
+    callee = overload->choice.getDeclOrNull();
     rawFnType = overload->openedType;
   } else {
     // If we didn't resolve an overload for the callee, we must be dealing with
     // a call of an arbitrary function expr.
     auto *call = cast<CallExpr>(anchor);
     assert(!shouldHaveDirectCalleeOverload(call) &&
            "Should we have resolved a callee for this?");
-    rawFnType = cs.getType(call->getFn())->getRValueType();
+    rawFnType = cs.getType(call->getFn());
   }
 
-  auto *fnType = resolveType(rawFnType)->getAs<FunctionType>();
+  // Try to resolve the function type by loading lvalues and looking through
+  // optional types, which can occur for expressions like `fn?(5)`.
+  auto *fnType = resolveType(rawFnType)
+                     ->getRValueType()
+                     ->lookThroughAllOptionalTypes()
+                     ->getAs<FunctionType>();
   if (!fnType)
     return None;
 

lib/Sema/ConstraintSystem.cpp

@@ -432,8 +432,15 @@ ConstraintLocator *ConstraintSystem::getCalleeLocator(Expr *expr) {
       return getConstraintLocator(fnLocator,
                                   ConstraintLocator::ConstructorMember);
     }
-    // Otherwise fall through and look for locators anchored on the fn expr.
-    expr = fnExpr->getSemanticsProvidingExpr();
+
+    // 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();
+    } else {
+      expr = fnExpr;
+    }
   }
 
   if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {

test/Constraints/optional.swift

@@ -10,8 +10,8 @@ class A {
   @objc(do_b_2:) func do_b(_ x: Int) {}
   @objc func do_b(_ x: Float) {}
 
-  @objc func do_c(x: Int) {}
-  @objc func do_c(y: Int) {}
+  @objc func do_c(x: Int) {} // expected-note {{incorrect labels for candidate (have: '(_:)', expected: '(x:)')}}
+  @objc func do_c(y: Int) {} // expected-note {{incorrect labels for candidate (have: '(_:)', expected: '(y:)')}}
 }
 
 func test0(_ a: AnyObject) {
@@ -20,7 +20,7 @@ func test0(_ a: AnyObject) {
   a.do_b?(1)
   a.do_b?(5.0)
 
-  a.do_c?(1) // expected-error {{cannot invoke value of function type with argument list '(Int)'}}
+  a.do_c?(1) // expected-error {{no exact matches in call to instance method 'do_c'}}
   a.do_c?(x: 1)
 }