[AST] Clean up handling of single-expression closures

include/swift/AST/Expr.h

@@ -3774,7 +3774,7 @@ class ClosureExpr : public AbstractClosureExpr {
   /// The bit indicates whether this closure has had a function builder
   /// applied to it.
   llvm::PointerIntPair<VarDecl *, 1, bool> CapturedSelfDeclAndAppliedBuilder;
-  
+
   /// The location of the "throws", if present.
   SourceLoc ThrowsLoc;
 
@@ -3786,7 +3786,8 @@ class ClosureExpr : public AbstractClosureExpr {
   SourceLoc InLoc;
 
   /// The explicitly-specified result type.
-  TypeExpr *ExplicitResultType;
+  llvm::PointerIntPair<TypeExpr *, 1, bool>
+    ExplicitResultTypeAndEnclosingChecked;
 
   /// The body of the closure, along with a bit indicating whether it
   /// was originally just a single expression.
@@ -3801,7 +3802,8 @@ class ClosureExpr : public AbstractClosureExpr {
       BracketRange(bracketRange),
       CapturedSelfDeclAndAppliedBuilder(capturedSelfDecl, false),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
-      ExplicitResultType(explicitResultType), Body(nullptr) {
+      ExplicitResultTypeAndEnclosingChecked(explicitResultType, false),
+      Body(nullptr) {
     setParameterList(params);
     Bits.ClosureExpr.HasAnonymousClosureVars = false;
   }
@@ -3855,13 +3857,15 @@ class ClosureExpr : public AbstractClosureExpr {
 
   Type getExplicitResultType() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultType->getInstanceType();
+    return ExplicitResultTypeAndEnclosingChecked.getPointer()
+        ->getInstanceType();
   }
   void setExplicitResultType(Type ty);
 
   TypeRepr *getExplicitResultTypeRepr() const {
     assert(hasExplicitResultType() && "No explicit result type");
-    return ExplicitResultType->getTypeRepr();
+    return ExplicitResultTypeAndEnclosingChecked.getPointer()
+        ->getTypeRepr();
   }
 
   /// Determine whether the closure has a single expression for its
@@ -3918,6 +3922,16 @@ class ClosureExpr : public AbstractClosureExpr {
     CapturedSelfDeclAndAppliedBuilder.setInt(flag);
   }
 
+  /// Whether this closure's body was type checked within the enclosing
+  /// context.
+  bool wasTypeCheckedInEnclosingContext() const {
+    return ExplicitResultTypeAndEnclosingChecked.getInt();
+  }
+
+  void setTypeCheckedInEnclosingContext(bool flag = true) {
+    ExplicitResultTypeAndEnclosingChecked.setInt(flag);
+  }
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Closure;
   }

lib/AST/ASTDumper.cpp

@@ -2496,11 +2496,7 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
     }
 
     OS << '\n';
-    if (E->hasSingleExpressionBody()) {
-      printRec(E->getSingleExpressionBody());
-    } else {
-      printRec(E->getBody(), E->getASTContext());
-    }
+    printRec(E->getBody(), E->getASTContext());
     PrintWithColorRAII(OS, ParenthesisColor) << ')';
   }
   void visitAutoClosureExpr(AutoClosureExpr *E) {

lib/AST/Expr.cpp

@@ -1923,7 +1923,8 @@ bool ClosureExpr::capturesSelfEnablingImplictSelf() const {
 
 void ClosureExpr::setExplicitResultType(Type ty) {
   assert(ty && !ty->hasTypeVariable());
-  ExplicitResultType->setType(MetatypeType::get(ty));
+  ExplicitResultTypeAndEnclosingChecked.getPointer()
+      ->setType(MetatypeType::get(ty));
 }
 
 FORWARD_SOURCE_LOCS_TO(AutoClosureExpr, Body)
@@ -2364,11 +2365,7 @@ void swift::simple_display(llvm::raw_ostream &out, const ClosureExpr *CE) {
     return;
   }
 
-  if (CE->hasSingleExpressionBody()) {
-    out << "single expression closure";
-  } else {
-    out << "closure";
-  }
+  out << "closure";
 }
 
 void swift::simple_display(llvm::raw_ostream &out,

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

@@ -331,18 +331,23 @@ SolutionApplicationToFunctionResult ConstraintSystem::applySolution(
     fn.setBody(newBody, /*isSingleExpression=*/false);
     if (closure) {
       closure->setAppliedFunctionBuilder();
+      closure->setTypeCheckedInEnclosingContext();
       solution.setExprTypes(closure);
     }
 
     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());
 
+    closure->setTypeCheckedInEnclosingContext();
     return SolutionApplicationToFunctionResult::Success;
   }
 

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/MiscDiagnostics.cpp

@@ -86,7 +86,7 @@ static void diagSyntacticUseRestrictions(const Expr *E, const DeclContext *DC,
 
     bool walkToDeclPre(Decl *D) override {
       if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
-        return closure->hasAppliedFunctionBuilder();
+        return closure->wasTypeCheckedInEnclosingContext();
       return false;
     }
 
@@ -1459,7 +1459,7 @@ static void diagnoseImplicitSelfUseInClosure(const Expr *E,
     // Don't walk into nested decls.
     bool walkToDeclPre(Decl *D) override {
       if (auto *closure = dyn_cast<ClosureExpr>(D->getDeclContext()))
-        return closure->hasAppliedFunctionBuilder();
+        return closure->wasTypeCheckedInEnclosingContext();
       return false;
     }
 

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();
+}

lib/Sema/TypeCheckStmt.cpp

@@ -140,10 +140,10 @@ namespace {
           }
         }
 
-        // If the closure has a single expression body or has had a function
-        // builder applied to it, we need to walk into it with a new sequence.
+        // If the closure was type checked within its enclosing context,
+        // we need to walk into it with a new sequence.
         // Otherwise, it'll have been separately type-checked.
-        if (CE->hasSingleExpressionBody() || CE->hasAppliedFunctionBuilder())
+        if (CE->wasTypeCheckedInEnclosingContext())
           CE->getBody()->walk(ContextualizeClosures(CE));
 
         TypeChecker::computeCaptures(CE);