[CS] Better handle unsupported result builder elements for completion

lib/Sema/BuilderTransform.cpp

@@ -269,12 +269,7 @@ class ResultBuilderTransform
     for (auto element : braceStmt->getElements()) {
       if (auto unsupported =
               transformBraceElement(element, newBody, buildBlockArguments)) {
-        // When in code completion mode, simply ignore unsported constructs to
-        // get results for anything that's unrelated to the unsupported
-        // constructs.
-        if (!ctx.CompletionCallback) {
-          return failTransform(*unsupported);
-        }
+        return failTransform(*unsupported);
       }
     }
 
@@ -780,16 +775,17 @@ class ResultBuilderTransform
 #undef UNSUPPORTED_STMT
 
 private:
-  static bool isBuildableIfChainRecursive(IfStmt *ifStmt, unsigned &numPayloads,
-                                          bool &isOptional) {
+  static void checkBuildableIfChainRecursive(IfStmt *ifStmt,
+                                             unsigned &numPayloads,
+                                             bool &isOptional) {
     // The 'then' clause contributes a payload.
     ++numPayloads;
 
     // If there's an 'else' clause, it contributes payloads:
     if (auto elseStmt = ifStmt->getElseStmt()) {
       // If it's 'else if', it contributes payloads recursively.
       if (auto elseIfStmt = dyn_cast<IfStmt>(elseStmt)) {
-        return isBuildableIfChainRecursive(elseIfStmt, numPayloads, isOptional);
+        checkBuildableIfChainRecursive(elseIfStmt, numPayloads, isOptional);
         // Otherwise it's just the one.
       } else {
         ++numPayloads;
@@ -799,8 +795,6 @@ class ResultBuilderTransform
     } else {
       isOptional = true;
     }
-
-    return true;
   }
 
   static bool hasUnconditionalElse(IfStmt *ifStmt) {
@@ -815,8 +809,7 @@ class ResultBuilderTransform
 
   bool isBuildableIfChain(IfStmt *ifStmt, unsigned &numPayloads,
                           bool &isOptional) {
-    if (!isBuildableIfChainRecursive(ifStmt, numPayloads, isOptional))
-      return false;
+    checkBuildableIfChainRecursive(ifStmt, numPayloads, isOptional);
 
     // If there's a missing 'else', we need 'buildOptional' to exist.
     if (isOptional && !builder.supportsOptional())
@@ -991,11 +984,11 @@ TypeChecker::applyResultBuilderBodyTransform(FuncDecl *func, Type builderType) {
   // of this decl; it's not part of the interface type.
   builderType = func->mapTypeIntoContext(builderType);
 
-  if (auto result = cs.matchResultBuilder(
-          func, builderType, resultContextType, resultConstraintKind,
-          /*contextualType=*/Type(),
-          cs.getConstraintLocator(func->getBody()))) {
-    if (result->isFailure())
+  {
+    auto result = cs.matchResultBuilder(
+        func, builderType, resultContextType, resultConstraintKind,
+        /*contextualType=*/Type(), cs.getConstraintLocator(func->getBody()));
+    if (!result || result->isFailure())
       return nullptr;
   }
 
@@ -1132,8 +1125,9 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
   if (fn.bodyHasExplicitReturnStmt()) {
     // Diagnostic mode means that solver couldn't reach any viable
     // solution, so let's diagnose presence of a `return` statement
-    // in the closure body.
-    if (shouldAttemptFixes()) {
+    // in the closure body. Avoid doing this for completion since we need to
+    // continue solving the body.
+    if (shouldAttemptFixes() && !isForCodeCompletion()) {
       if (recordFix(IgnoreResultBuilderWithReturnStmts::create(
               *this, builderType,
               getConstraintLocator(fn.getAbstractClosureExpr()))))
@@ -1157,18 +1151,17 @@ ConstraintSystem::matchResultBuilder(AnyFunctionRef fn, Type builderType,
     auto *body = transform.apply(fn.getBody());
 
     if (auto unsupported = transform.getUnsupportedElement()) {
-      assert(!body || getASTContext().CompletionCallback);
-
-      // If we aren't supposed to attempt fixes, fail.
-      if (!shouldAttemptFixes()) {
-        return getTypeMatchFailure(locator);
-      }
+      assert(!body);
 
       // If we're solving for code completion and the body contains the code
-      // completion location, skipping it won't get us to a useful solution so
-      // just bail.
+      // completion location, fall back to solving as a regular function body.
       if (isForCodeCompletion() &&
           containsIDEInspectionTarget(fn.getBody())) {
+        return std::nullopt;
+      }
+
+      // If we aren't supposed to attempt fixes, fail.
+      if (!shouldAttemptFixes()) {
         return getTypeMatchFailure(locator);
       }
 

test/IDE/complete_in_result_builder.swift

@@ -16,7 +16,10 @@ struct TupleBuilder<T> {
   }
 }
 
-func buildStringTuple<Result>(@TupleBuilder<String> x: () -> Result) {}
+@discardableResult
+func buildStringTuple<Result>(@TupleBuilder<String> x: () -> Result) -> Result {
+  x()
+}
 
 enum StringFactory {
   static func makeString(x: String) -> String { return x }
@@ -50,6 +53,12 @@ func testGlobalLookup() {
     }
   }
 
+  buildStringTuple {
+    if {
+      #^GLOBAL_LOOKUP_IN_IF_BODY_WITHOUT_CONDITION_CLOSURE?check=GLOBAL_LOOKUP_IN_IF_BODY^#
+    }
+  }
+
   @TupleBuilder<String> var x4 {
     guard else {
       #^GLOBAL_LOOKUP_IN_GUARD_BODY_WITHOUT_CONDITION?check=GLOBAL_LOOKUP_IN_IF_BODY^#
@@ -98,6 +107,21 @@ struct FooStruct {
   func intGen() -> Int { return 1 }
 }
 
+func testReturn() {
+  // Make sure we can still complete even with return.
+  buildStringTuple {
+    StringFactory.#^COMPLETE_STATIC_MEMBER_WITH_RETURN^#
+    // COMPLETE_STATIC_MEMBER_WITH_RETURN: Decl[StaticMethod]/CurrNominal:     makeString({#x: String#})[#String#];
+    return ""
+  }
+
+  buildStringTuple {
+    ""
+    return FooStruct()
+  }.#^COMPLETE_INSTANCE_MEMBER_ON_RETURN_BUILDER^#
+  // COMPLETE_INSTANCE_MEMBER_ON_RETURN_BUILDER: Decl[InstanceVar]/CurrNominal: instanceVar[#Int#]; name=instanceVar
+}
+
 func testPatternMatching() {
   @TupleBuilder<String> var x1 {
     let x = Letters.b
@@ -128,6 +152,11 @@ func testPatternMatching() {
     let x: FooStruct? = FooStruct()
     guard case .#^GUARD_CASE_PATTERN_2?check=OPTIONAL_FOOSTRUCT^#some() = x {}
   }
+
+  buildStringTuple {
+    let x: FooStruct? = FooStruct()
+    guard case .#^GUARD_CASE_PATTERN_3?check=OPTIONAL_FOOSTRUCT^#some() = x {}
+  }
 }
 
 func testCompleteFunctionArgumentLabels() {
@@ -410,3 +439,24 @@ func testInForLoop(_ x: [Int]) {
     }
   }
 }
+
+func testUnsupportedForLoop() {
+  @resultBuilder
+  struct Builder {
+    static func buildBlock<T>(_ components: T...) -> T {
+      components.first!
+    }
+  }
+
+  func foo(@Builder fn: () -> Int) {}
+
+  foo {
+    for i in [0].#^COMPLETE_IN_UNSUPPORTED_FOR^# {}
+  }
+  // COMPLETE_IN_UNSUPPORTED_FOR: Decl[InstanceVar]/Super/IsSystem: first[#Int?#]; name=first
+
+  foo {
+    for i in [0].#^COMPLETE_IN_UNSUPPORTED_FOR_2?check=COMPLETE_IN_UNSUPPORTED_FOR^# {}
+    return 0
+  }
+}