[Sema] Type check multi-statement closures inside result builders

Author: ahoppen

include/swift/Sema/ConstraintSystem.h

@@ -73,8 +73,9 @@ class SolutionApplicationTarget;
 // so they could be made friends of ConstraintSystem.
 namespace TypeChecker {
 
-Optional<BraceStmt *> applyResultBuilderBodyTransform(FuncDecl *func,
-                                                        Type builderType);
+Optional<BraceStmt *> applyResultBuilderBodyTransform(
+    FuncDecl *func, Type builderType,
+    bool ClosuresInResultBuilderDontParticipateInInference);
 
 Optional<constraints::SolutionApplicationTarget>
 typeCheckExpression(constraints::SolutionApplicationTarget &target,
@@ -1774,6 +1775,12 @@ enum class ConstraintSystemFlags {
 
   /// When set, ignore async/sync mismatches
   IgnoreAsyncSyncMismatch = 0x80,
+
+  /// Non solver-based code completion expects that closures inside result
+  /// builders don't participate in inference.
+  /// Once all code completion kinds are migrated to solver-based we should be
+  /// able to remove this flag.
+  ClosuresInResultBuildersDontParticipateInInference = 0x100,
 };
 
 /// Options that affect the constraint system as a whole.
@@ -3576,8 +3583,9 @@ class ConstraintSystem {
 
   // FIXME: Perhaps these belong on ConstraintSystem itself.
   friend Optional<BraceStmt *>
-  swift::TypeChecker::applyResultBuilderBodyTransform(FuncDecl *func,
-                                                        Type builderType);
+  swift::TypeChecker::applyResultBuilderBodyTransform(
+      FuncDecl *func, Type builderType,
+      bool ClosuresInResultBuilderDontParticipateInInference);
 
   friend Optional<SolutionApplicationTarget>
   swift::TypeChecker::typeCheckExpression(

lib/Sema/BuilderTransform.cpp

@@ -922,7 +922,8 @@ class ResultBuilderTransform
 } // end anonymous namespace
 
 Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
-    FuncDecl *func, Type builderType) {
+    FuncDecl *func, Type builderType,
+    bool ClosuresInResultBuilderDontParticipateInInference) {
   // Pre-check the body: pre-check any expressions in it and look
   // for return statements.
   //
@@ -978,6 +979,10 @@ Optional<BraceStmt *> TypeChecker::applyResultBuilderBodyTransform(
   }
 
   ConstraintSystemOptions options = ConstraintSystemFlags::AllowFixes;
+  if (ClosuresInResultBuilderDontParticipateInInference) {
+    options |= ConstraintSystemFlags::
+        ClosuresInResultBuildersDontParticipateInInference;
+  }
   auto resultInterfaceTy = func->getResultInterfaceType();
   auto resultContextType = func->mapTypeIntoContext(resultInterfaceTy);
 

lib/Sema/CSDiagnostics.cpp

@@ -4874,7 +4874,9 @@ bool MissingArgumentsFailure::diagnoseClosure(const ClosureExpr *closure) {
 
   auto *locator = getLocator();
   if (locator->isForContextualType()) {
-    funcType = getContextualType(locator->getAnchor())->getAs<FunctionType>();
+    if (auto type = getContextualType(locator->getAnchor())) {
+      funcType = type->getAs<FunctionType>();
+    }
   } else if (auto info = getFunctionArgApplyInfo(locator)) {
     auto paramType = info->getParamType();
     // Drop a single layer of optionality because argument could get injected

lib/Sema/CSFix.cpp

@@ -602,6 +602,10 @@ bool AllowTupleTypeMismatch::coalesceAndDiagnose(
     purpose = cs.getContextualTypePurpose(locator->getAnchor());
   }
 
+  if (!getFromType()->is<TupleType>() || !getToType()->is<TupleType>()) {
+    return false;
+  }
+
   TupleContextualFailure failure(solution, purpose, getFromType(), getToType(),
                                  indices, locator);
   return failure.diagnose(asNote);

lib/Sema/ConstraintSystem.cpp

@@ -6892,7 +6892,13 @@ bool ConstraintSystem::participatesInInference(ClosureExpr *closure) const {
 
   // If body is nested in a parent that has a function builder applied,
   // let's prevent inference until result builders.
-  return !isInResultBuilderContext(closure);
+  if (Options.contains(
+          ConstraintSystemFlags::
+              ClosuresInResultBuildersDontParticipateInInference)) {
+    return !isInResultBuilderContext(closure);
+  } else {
+    return true;
+  }
 }
 
 TypeVarBindingProducer::TypeVarBindingProducer(BindingSet &bindings)

lib/Sema/TypeCheckStmt.cpp

@@ -1946,8 +1946,10 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
   // Function builder function doesn't support partial type checking.
   if (auto *func = dyn_cast<FuncDecl>(DC)) {
     if (Type builderType = getResultBuilderType(func)) {
-      auto optBody =
-          TypeChecker::applyResultBuilderBodyTransform(func, builderType);
+      auto optBody = TypeChecker::applyResultBuilderBodyTransform(
+          func, builderType,
+          /*ClosuresInResultBuilderDontParticipateInInference=*/
+              ctx.CompletionCallback == nullptr);
       if (optBody && *optBody) {
         // Wire up the function body now.
         func->setBody(*optBody, AbstractFunctionDecl::BodyKind::TypeChecked);

lib/Sema/TypeChecker.h

@@ -451,8 +451,9 @@ void typeCheckASTNode(ASTNode &node, DeclContext *DC,
 /// e.g., because of a \c return statement. Otherwise, returns either the
 /// fully type-checked body of the function (on success) or a \c nullptr
 /// value if an error occurred while type checking the transformed body.
-Optional<BraceStmt *> applyResultBuilderBodyTransform(FuncDecl *func,
-                                                        Type builderType);
+Optional<BraceStmt *> applyResultBuilderBodyTransform(
+    FuncDecl *func, Type builderType,
+    bool ClosuresInResultBuilderDontParticipateInInference = true);
 
 /// Find the return statements within the body of the given function.
 std::vector<ReturnStmt *> findReturnStatements(AnyFunctionRef fn);

test/Constraints/result_builder_diags.swift

@@ -78,7 +78,7 @@ struct TupleBuilderWithoutIf { // expected-note 3{{struct 'TupleBuilderWithoutIf
   static func buildDo<T>(_ value: T) -> T { return value }
 }
 
-func tuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) {
+func tuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) { // expected-note 3 {{in call to function 'tuplify(_:body:)'}}
   print(body(cond))
 }
 
@@ -88,7 +88,7 @@ func tuplifyWithoutIf<T>(_ cond: Bool, @TupleBuilderWithoutIf body: (Bool) -> T)
 
 func testDiags() {
   // For loop
-  tuplify(true) { _ in
+  tuplify(true) { _ in // expected-error {{generic parameter 'T' could not be inferred}}
     17
     for c in name {
     // expected-error@-1 {{cannot find 'name' in scope}}
@@ -466,7 +466,7 @@ struct TestConstraintGenerationErrors {
   }
 
   func buildTupleClosure() {
-    tuplify(true) { _ in
+    tuplify(true) { _ in // expected-error {{generic parameter 'T' could not be inferred}}
       let a = nothing // expected-error {{cannot find 'nothing' in scope}}
       String(nothing) // expected-error {{cannot find 'nothing' in scope}}
     }
@@ -733,7 +733,7 @@ struct TuplifiedStructWithInvalidClosure {
   }
 
   @TupleBuilder var nestedErrorsDiagnosedByParser: some Any {
-    tuplify(true) { _ in
+    tuplify(true) { _ in // expected-error {{generic parameter 'T' could not be inferred}}
       tuplify { _ in
         self. // expected-error {{expected member name following '.'}}
       }

test/Parse/recovery.swift

@@ -773,7 +773,7 @@ class r22240342 {
   lazy var xx: Int = {
     foo {  // expected-error {{cannot find 'foo' in scope}}
       let issueView = 42
-      issueView.delegate = 12
+      issueView.delegate = 12 // expected-error {{value of type 'Int' has no member 'delegate'}}
 
     }
     return 42

test/StringProcessing/Sema/regex_builder_fix_import_top_level.swift

@@ -19,11 +19,7 @@ Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be
   /g(h)(i)/
 }
 
-// FIXME: Unfortunately we bail from CSGen if we end up with an ErrorExpr, so
-// don't get a chance to diagnose. We ought to try solving with holes.
-// For now at least, this error should at least hopefully nudge users into
-// realizing they have a missing import.
-Regex {
+Regex { // expected-error {{regex builder requires the 'RegexBuilder' module be imported'}}
   Capture { // expected-error {{cannot find 'Capture' in scope}}
     /abc/
   }