[Constaint system] Add a predicate for when to check closures in enclosing expr

Introduce a new predicate, shouldTypeCheckInEnclosingExpression(), to
determine when the body of a closure should be checked as part of the
enclosing expression rather than separately, and use it in the various
places where "hasSingleExpressionBody()" was used for that purpose.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -1701,6 +1701,13 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
 FunctionBuilderBodyPreCheck
 PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval,
                                          AnyFunctionRef fn) const {
+  // We don't want to do the precheck if it will already have happened in
+  // the enclosing expression.
+  bool skipPrecheck = false;
+  if (auto closure = dyn_cast_or_null<ClosureExpr>(
+          fn.getAbstractClosureExpr()))
+    skipPrecheck = shouldTypeCheckInEnclosingExpression(closure);
+
   return PreCheckFunctionBuilderApplication(fn, false).run();
 }
 

lib/Sema/CSApply.cpp

@@ -5787,9 +5787,9 @@ static bool applyTypeToClosureExpr(ConstraintSystem &cs,
   if (auto CE = dyn_cast<ClosureExpr>(expr)) {
     cs.setType(CE, toType);
 
-    // If this is not a single-expression closure, write the type into the
-    // ClosureExpr directly here, since the visitor won't.
-    if (!CE->hasSingleExpressionBody())
+    // If this closure isn't type-checked in its enclosing expression, write
+    // theg type into the ClosureExpr directly here, since the visitor won't.
+    if (!shouldTypeCheckInEnclosingExpression(CE))
       CE->setType(toType);
 
     return true;

lib/Sema/CSClosure.cpp

@@ -336,10 +336,13 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
     }
 
     return SolutionApplicationToFunctionResult::Success;
+
   }
+  assert(closure && "Can only get here with a closure at the moment");
 
-  // If there is a single-expression body, transform that body now.
-  if (fn.hasSingleExpressionBody()) {
+  // If this closure is checked as part of the enclosing expression, handle
+  // that now.
+  if (shouldTypeCheckInEnclosingExpression(closure)) {
     ClosureConstraintApplication application(
         solution, closure, closureFnType->getResult(), rewriteTarget);
     application.visit(fn.getBody());

lib/Sema/CSGen.cpp

@@ -2248,7 +2248,8 @@ namespace {
           // as potential hole right away.
           resultTy = CS.createTypeVariable(
               resultLoc,
-              closure->hasSingleExpressionBody() ? 0 : TVO_CanBindToHole);
+              shouldTypeCheckInEnclosingExpression(closure)
+                ? 0 : TVO_CanBindToHole);
         }
       }
 

lib/Sema/CSSimplify.cpp

@@ -7157,14 +7157,13 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
     }
   }
 
-  bool hasReturn = hasExplicitResult(closure);
+  // If this closure should be type-checked as part of this expression,
+  // generate constraints for it now.
   auto &ctx = getASTContext();
-  // If this is a multi-statement closure its body doesn't participate
-  // in type-checking.
-  if (closure->hasSingleExpressionBody()) {
+  if (shouldTypeCheckInEnclosingExpression(closure)) {
     if (generateConstraints(closure, closureType->getResult()))
       return false;
-  } else if (!hasReturn) {
+  } else if (!hasExplicitResult(closure)) {
     // If this closure has an empty body and no explicit result type
     // let's bind result type to `Void` since that's the only type empty body
     // can produce. Otherwise, if (multi-statement) closure doesn't have

lib/Sema/ConstraintSystem.h

@@ -5487,6 +5487,12 @@ BraceStmt *applyFunctionBuilderTransform(
           constraints::SolutionApplicationTarget)>
             rewriteTarget);
 
+/// Determine whether the given closure expression should be type-checked
+/// within the context of its enclosing expression. Otherwise, it will be
+/// separately type-checked once its enclosing expression has determined all
+/// of the parameter and result types without looking at the body.
+bool shouldTypeCheckInEnclosingExpression(ClosureExpr *expr);
+
 } // end namespace swift
 
 #endif // LLVM_SWIFT_SEMA_CONSTRAINT_SYSTEM_H

lib/Sema/TypeCheckAvailability.cpp

@@ -1128,7 +1128,8 @@ static void findAvailabilityFixItNodes(SourceRange ReferenceRange,
       [](ASTNode Node, ASTWalker::ParentTy Parent) {
         if (Expr *ParentExpr = Parent.getAsExpr()) {
           auto *ParentClosure = dyn_cast<ClosureExpr>(ParentExpr);
-          if (!ParentClosure || !ParentClosure->hasSingleExpressionBody()) {
+          if (!ParentClosure ||
+              !ParentClosure->wasTypeCheckedInEnclosingContext()) {
             return false;
           }
         } else if (auto *ParentStmt = Parent.getAsStmt()) {

lib/Sema/TypeCheckConstraints.cpp

@@ -1320,9 +1320,8 @@ bool PreCheckExpression::walkToClosureExprPre(ClosureExpr *closure) {
   if (hadParameterError)
     return false;
 
-  // If the closure has a multi-statement body, we don't walk into it
-  // here.
-  if (!closure->hasSingleExpressionBody())
+  // If we won't be checking the body of the closure, don't walk into it here.
+  if (!shouldTypeCheckInEnclosingExpression(closure))
     return false;
 
   // Update the current DeclContext to be the closure we're about to
@@ -4058,3 +4057,7 @@ HasDynamicCallableAttributeRequest::evaluate(Evaluator &evaluator,
     return type->hasDynamicCallableAttribute();
   });
 }
+
+bool swift::shouldTypeCheckInEnclosingExpression(ClosureExpr *expr) {
+  return expr->hasSingleExpressionBody();
+}