[CS] Simplify getCalleeDeclAndArgs

Now that we associate argument labels for key path
subscript components, remove the special logic for
it. In addition, we can now search for callees
directly by using `getCalleeLocator`, as it should
now be able to find all the correct callees that
`getCalleeDeclAndArgs` does.

By using `getCalleeLocator`, we now also correctly
resolve callees for operator calls, meaning that
we can now use them with function builders. In
addition, we no longer incorrectly resolve callees
for calls made on SubscriptExprs.

Resolves SR-11439 & SR-11440.

Author: hamishknight

lib/Sema/CSSimplify.cpp

@@ -735,124 +735,54 @@ static std::tuple<ValueDecl *, bool, ArrayRef<Identifier>, bool,
                   ConstraintLocator *>
 getCalleeDeclAndArgs(ConstraintSystem &cs,
                      ConstraintLocatorBuilder callBuilder) {
-  ArrayRef<Identifier> argLabels;
-  bool hasTrailingClosure = false;
-  ConstraintLocator *targetLocator = nullptr;
+  auto formUnknownCallee =
+      []() -> std::tuple<ValueDecl *, bool, ArrayRef<Identifier>, bool,
+                         ConstraintLocator *> {
+    return std::make_tuple(/*decl*/ nullptr, /*hasAppliedSelf*/ false,
+                           /*argLabels*/ ArrayRef<Identifier>(),
+                           /*hasTrailingClosure*/ false,
+                           /*calleeLocator*/ nullptr);
+  };
 
   auto *callLocator = cs.getConstraintLocator(callBuilder);
   auto *callExpr = callLocator->getAnchor();
 
   // Break down the call.
   if (!callExpr)
-    return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                           hasTrailingClosure, targetLocator);
+    return formUnknownCallee();
 
-  auto path = callLocator->getPath();
   // Our remaining path can only be 'ApplyArgument'.
+  auto path = callLocator->getPath();
   if (!path.empty() &&
       !(path.size() <= 2 &&
         path.back().getKind() == ConstraintLocator::ApplyArgument))
-    return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                           hasTrailingClosure, targetLocator);
+    return formUnknownCallee();
 
   // Dig out the callee information.
-  if (auto argInfo = cs.getArgumentInfo(callLocator)) {
-    argLabels = argInfo->Labels;
-    hasTrailingClosure = argInfo->HasTrailingClosure;
-    targetLocator = cs.getConstraintLocator(
-        isa<CallExpr>(callExpr) ? cast<CallExpr>(callExpr)->getDirectCallee()
-                                : callExpr);
-  } else if (auto keyPath = dyn_cast<KeyPathExpr>(callExpr)) {
-    if (path.size() != 2)
-      return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                             hasTrailingClosure, nullptr);
-
-    // We must have a KeyPathComponent followed by an ApplyArgument.
-    auto componentElt = path[0].getAs<LocatorPathElt::KeyPathComponent>();
-    if (!componentElt || path[1].getKind() != ConstraintLocator::ApplyArgument)
-      return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                             hasTrailingClosure, nullptr);
-
-    auto componentIndex = componentElt->getIndex();
-    if (componentIndex >= keyPath->getComponents().size())
-      return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                             hasTrailingClosure, nullptr);
-
-    auto &component = keyPath->getComponents()[componentIndex];
-    switch (component.getKind()) {
-    case KeyPathExpr::Component::Kind::Subscript:
-    case KeyPathExpr::Component::Kind::UnresolvedSubscript:
-      targetLocator = cs.getConstraintLocator(callExpr, path[0]);
-      argLabels = component.getSubscriptLabels();
-      hasTrailingClosure = false; // key paths don't support trailing closures
-      break;
-      
-    case KeyPathExpr::Component::Kind::Invalid:
-    case KeyPathExpr::Component::Kind::UnresolvedProperty:
-    case KeyPathExpr::Component::Kind::Property:
-    case KeyPathExpr::Component::Kind::OptionalForce:
-    case KeyPathExpr::Component::Kind::OptionalChain:
-    case KeyPathExpr::Component::Kind::OptionalWrap:
-    case KeyPathExpr::Component::Kind::Identity:
-    case KeyPathExpr::Component::Kind::TupleElement:
-      return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                             hasTrailingClosure, nullptr);
-    }
-  } else {
-    return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                           hasTrailingClosure, targetLocator);
-  }
+  auto argInfo = cs.getArgumentInfo(callLocator);
+  if (!argInfo)
+    return formUnknownCallee();
+
+  auto argLabels = argInfo->Labels;
+  auto hasTrailingClosure = argInfo->HasTrailingClosure;
+  auto calleeLocator = cs.getCalleeLocator(callLocator);
 
   // Find the overload choice corresponding to the callee locator.
   // FIXME: This linearly walks the list of resolved overloads, which is
   // potentially very expensive.
-  Optional<OverloadChoice> choice;
-  ConstraintLocator *calleeLocator = nullptr;
-  for (auto resolved = cs.getResolvedOverloadSets(); resolved;
-       resolved = resolved->Previous) {
-    // FIXME: Workaround null locators.
-    if (!resolved->Locator) continue;
-
-    auto resolvedLocator = resolved->Locator;
-    SmallVector<LocatorPathElt, 4> resolvedPath(
-                                     resolvedLocator->getPath().begin(),
-                                     resolvedLocator->getPath().end());
-    if (!resolvedPath.empty() &&
-        (resolvedPath.back().getKind() == ConstraintLocator::SubscriptMember ||
-         resolvedPath.back().getKind() == ConstraintLocator::Member ||
-         resolvedPath.back().getKind() == ConstraintLocator::UnresolvedMember ||
-         resolvedPath.back().getKind() ==
-           ConstraintLocator::ConstructorMember)) {
-      resolvedPath.pop_back();
-      resolvedLocator = cs.getConstraintLocator(
-                          resolvedLocator->getAnchor(),
-                          resolvedPath,
-                          resolvedLocator->getSummaryFlags());
-    }
-
-    SourceRange range;
-    resolvedLocator = simplifyLocator(cs, resolvedLocator, range);
-
-    if (resolvedLocator == targetLocator) {
-      calleeLocator = resolved->Locator;
-      choice = resolved->Choice;
-      break;
-    }
-  }
-  
-  // If we didn't find any matching overloads, we're done.
-  if (!choice)
-    return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                           hasTrailingClosure, nullptr);
-
-  // If there's a declaration, return it.
-  if (auto *decl = choice->getDeclOrNull()) {
-    return std::make_tuple(decl, hasAppliedSelf(cs, *choice), argLabels,
-                           hasTrailingClosure, calleeLocator);
-  }
-
-  return std::make_tuple(nullptr, /*hasAppliedSelf=*/false, argLabels,
-                         hasTrailingClosure, calleeLocator);
+  auto selectedOverload = cs.findSelectedOverloadFor(calleeLocator);
+
+  // If we didn't find any matching overloads, we're done. Just return the
+  // argument info.
+  if (!selectedOverload)
+    return std::make_tuple(/*decl*/ nullptr, /*hasAppliedSelf*/ false,
+                           argLabels, hasTrailingClosure,
+                           /*calleeLocator*/ nullptr);
+
+  // Return the found declaration, assuming there is one.
+  auto choice = selectedOverload->Choice;
+  return std::make_tuple(choice.getDeclOrNull(), hasAppliedSelf(cs, choice),
+                         argLabels, hasTrailingClosure, calleeLocator);
 }
 
 class ArgumentFailureTracker : public MatchCallArgumentListener {

test/Constraints/function_builder.swift

@@ -212,6 +212,18 @@ print(mbuilders.methodBuilder(13))
 // CHECK: ("propertyBuilder", 12)
 print(mbuilders.propertyBuilder)
 
+// SR-11439: Operator builders
+infix operator ^^^
+func ^^^ (lhs: Int, @TupleBuilder rhs: (Int) -> (String, Int)) -> (String, Int) {
+  return rhs(lhs)
+}
+
+// CHECK: ("hello", 6)
+print(5 ^^^ {
+  "hello"
+  $0 + 1
+})
+
 struct Tagged<Tag, Entity> {
   let tag: Tag
   let entity: Entity

test/Constraints/function_builder_diags.swift

@@ -172,3 +172,30 @@ func test_51167632() -> some P {
     // expected-note@-1 {{explicitly specify the generic arguments to fix this issue}} {{10-10=<<#L: P#>>}}
   })
 }
+
+struct SR11440 {
+  typealias ReturnsTuple<T> = () -> (T, T)
+  subscript<T, U>(@TupleBuilder x: ReturnsTuple<T>) -> (ReturnsTuple<U>) -> Void { //expected-note {{in call to 'subscript(_:)'}}
+    return { _ in }
+  }
+
+  func foo() {
+    // This is okay, we apply the function builder for the subscript arg.
+    self[{
+      5
+      5
+    }]({
+      (5, 5)
+    })
+
+    // But we shouldn't perform the transform for the argument to the call
+    // made on the function returned from the subscript.
+    self[{ // expected-error {{generic parameter 'U' could not be inferred}}
+      5
+      5
+    }]({
+      5
+      5
+    })
+  }
+}