Merge pull request #74352 from hamishknight/complete-disjunction

[Completion] Skip `tryOptimizeGenericDisjunction` if there's a completion child

Author: hamishknight

include/swift/Sema/ConstraintSystem.h

@@ -2826,6 +2826,10 @@ class ConstraintSystem {
   /// Associate an argument list with a call at a given locator.
   void associateArgumentList(ConstraintLocator *locator, ArgumentList *args);
 
+  /// If the given node is a function expression with a parent ApplyExpr,
+  /// returns the apply, otherwise returns the node itself.
+  ASTNode includingParentApply(ASTNode node);
+
   std::optional<SelectedOverload>
   findSelectedOverloadFor(ConstraintLocator *locator) const {
     auto result = ResolvedOverloads.find(locator);

lib/Sema/CSRanking.cpp

@@ -1036,18 +1036,9 @@ SolutionCompareResult ConstraintSystem::compareSolutions(
     if (cs.isForCodeCompletion()) {
       // Don't rank based on overload choices of function calls that contain the
       // code completion token.
-      ASTNode anchor = simplifyLocatorToAnchor(overload.locator);
-      if (auto expr = getAsExpr(anchor)) {
-        // If the anchor is a called function, also don't rank overload choices
-        // if any of the arguments contain the code completion token.
-        if (auto apply = dyn_cast_or_null<ApplyExpr>(cs.getParentExpr(expr))) {
-          if (apply->getFn() == expr) {
-            anchor = apply;
-          }
-        }
-      }
-      if (anchor && cs.containsIDEInspectionTarget(anchor)) {
-        continue;
+      if (auto anchor = simplifyLocatorToAnchor(overload.locator)) {
+        if (cs.containsIDEInspectionTarget(cs.includingParentApply(anchor)))
+          continue;
       }
     }
 

lib/Sema/CSSolver.cpp

@@ -2064,9 +2064,20 @@ Constraint *ConstraintSystem::getUnboundBindOverloadDisjunction(
 
 // Performance hack: if there are two generic overloads, and one is
 // more specialized than the other, prefer the more-specialized one.
-static Constraint *tryOptimizeGenericDisjunction(
-                                          DeclContext *dc,
-                                          ArrayRef<Constraint *> constraints) {
+static Constraint *
+tryOptimizeGenericDisjunction(ConstraintSystem &cs, Constraint *disjunction,
+                              ArrayRef<Constraint *> constraints) {
+  auto *dc = cs.DC;
+
+  // If we're solving for code completion, and have a child completion token,
+  // skip this optimization since the completion token being a placeholder can
+  // allow us to prefer an unhelpful disjunction choice.
+  if (cs.isForCodeCompletion()) {
+    auto anchor = disjunction->getLocator()->getAnchor();
+    if (cs.containsIDEInspectionTarget(cs.includingParentApply(anchor)))
+      return nullptr;
+  }
+
   llvm::SmallVector<Constraint *, 4> choices;
   for (auto *choice : constraints) {
     if (choices.size() > 2)
@@ -2311,7 +2322,7 @@ void DisjunctionChoiceProducer::partitionDisjunction(
     SmallVectorImpl<unsigned> &PartitionBeginning) {
   // Apply a special-case rule for favoring one generic function over
   // another.
-  if (auto favored = tryOptimizeGenericDisjunction(CS.DC, Choices)) {
+  if (auto favored = tryOptimizeGenericDisjunction(CS, Disjunction, Choices)) {
     CS.favorConstraint(favored);
   }
 

lib/Sema/ConstraintSystem.cpp

@@ -6667,6 +6667,16 @@ static bool shouldHaveDirectCalleeOverload(const CallExpr *callExpr) {
 }
 #endif
 
+ASTNode ConstraintSystem::includingParentApply(ASTNode node) {
+  if (auto *expr = getAsExpr(node)) {
+    if (auto apply = getAsExpr<ApplyExpr>(getParentExpr(expr))) {
+      if (apply->getFn() == expr)
+        return apply;
+    }
+  }
+  return node;
+}
+
 Type Solution::resolveInterfaceType(Type type) const {
   auto resolvedType = type.transform([&](Type type) -> Type {
     if (auto *tvt = type->getAs<TypeVariableType>()) {

test/IDE/complete_rdar127844278.swift

@@ -0,0 +1,35 @@
+// RUN: %batch-code-completion
+
+// rdar://127844278 - Make sure we don't attempt to favor one buildExpression
+// call over the other when there's a child code completion.
+
+protocol P {}
+protocol Q: P {}
+
+struct S<T>: P {
+  private init() {}
+}
+
+extension S where T == () {
+  init(a: Void) {}
+}
+extension S: Q where T == String {
+  init(b: String) {}
+}
+
+@resultBuilder struct Builder {
+  static func buildExpression<T: P>(_ x: T) -> T { x }
+  static func buildExpression<T: Q>(_ x: T) -> T { x }
+
+  static func buildBlock<T: P>(_ x: T) -> some P { x }
+  static func buildBlock<T: Q>(_ x: T) -> some P { x }
+}
+
+func foo<T>(@Builder _: () -> T) where T: P {}
+func foo<T>(@Builder _: () -> T) where T: Q {}
+
+foo {
+  S(#^COMPLETE^#)
+}
+// COMPLETE-DAG: Decl[Constructor]/CurrNominal/Flair[ArgLabels]: ['(']{#a: Void#}[')'][#S<()>#]; name=a:
+// COMPLETE-DAG: Decl[Constructor]/CurrNominal/Flair[ArgLabels]: ['(']{#b: String#}[')'][#S<String>#]; name=b: