[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,
@@ -1801,6 +1802,12 @@ enum class ConstraintSystemFlags {
 
   /// Disable macro expansions.
   DisableMacroExpansions = 0x100,
+
+  /// 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 = 0x200,
 };
 
 /// Options that affect the constraint system as a whole.
@@ -3609,8 +3616,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

@@ -2224,7 +2224,8 @@ BraceStmt *swift::applyResultBuilderTransform(
 }
 
 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.
   //
@@ -2280,6 +2281,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/CSFix.cpp

@@ -622,6 +622,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

@@ -7143,7 +7143,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

@@ -2336,8 +2336,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

@@ -457,8 +457,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);