[Function builders] Make sure we contextualize closures with builders applied.

We used to get this contextualization "for free" because closures that
had function builders applied to them would get translated into
single-expression closures. Now, we need to check for this explicitly.

Author: DougGregor

include/swift/AST/Expr.h

@@ -3599,7 +3599,10 @@ class ClosureExpr : public AbstractClosureExpr {
   /// the CaptureListExpr which would normally maintain this sort of
   /// information about captured variables), we need to have some way to access
   /// this information directly on the ClosureExpr.
-  VarDecl *CapturedSelfDecl;
+  ///
+  /// 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;
@@ -3624,7 +3627,8 @@ class ClosureExpr : public AbstractClosureExpr {
               unsigned discriminator, DeclContext *parent)
     : AbstractClosureExpr(ExprKind::Closure, Type(), /*Implicit=*/false,
                           discriminator, parent),
-      BracketRange(bracketRange), CapturedSelfDecl(capturedSelfDecl),
+      BracketRange(bracketRange),
+      CapturedSelfDeclAndAppliedBuilder(capturedSelfDecl, false),
       ThrowsLoc(throwsLoc), ArrowLoc(arrowLoc), InLoc(inLoc),
       ExplicitResultType(explicitResultType), Body(nullptr) {
     setParameterList(params);
@@ -3726,13 +3730,23 @@ class ClosureExpr : public AbstractClosureExpr {
   bool hasEmptyBody() const;
 
   /// VarDecl captured by this closure under the literal name \c self , if any.
-  VarDecl *getCapturedSelfDecl() const { return CapturedSelfDecl; }
+  VarDecl *getCapturedSelfDecl() const {
+    return CapturedSelfDeclAndAppliedBuilder.getPointer();
+  }
 
   /// Whether this closure captures the \c self param in its body in such a
   /// way that implicit \c self is enabled within its body (i.e. \c self is
   /// captured non-weakly).
   bool capturesSelfEnablingImplictSelf() const;
 
+  bool hasAppliedFunctionBuilder() const {
+    return CapturedSelfDeclAndAppliedBuilder.getInt();
+  }
+
+  void setAppliedFunctionBuilder(bool flag = true) {
+    CapturedSelfDeclAndAppliedBuilder.setInt(flag);
+  }
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Closure;
   }

lib/Sema/CSApply.cpp

@@ -7074,6 +7074,7 @@ namespace {
                 return Rewriter.coerceToType(expr, toType, locator);
               });
           closure->setBody(newBody, /*isSingleExpression=*/false);
+          closure->setAppliedFunctionBuilder();
 
           Rewriter.solution.setExprTypes(closure);
         } else if (closure->hasSingleExpressionBody()) {

lib/Sema/TypeCheckStmt.cpp

@@ -144,10 +144,10 @@ namespace {
           }
         }
 
-        // If the closure has a single expression body, we need to walk into it
-        // with a new sequence.  Otherwise, it'll have been separately
-        // type-checked.
-        if (CE->hasSingleExpressionBody())
+        // 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.
+        // Otherwise, it'll have been separately type-checked.
+        if (CE->hasSingleExpressionBody() || CE->hasAppliedFunctionBuilder())
           CE->getBody()->walk(ContextualizeClosures(CE));
 
         TypeChecker::computeCaptures(CE);

test/Constraints/function_builder.swift

@@ -352,14 +352,20 @@ func acceptComponentBuilder(@ComponentBuilder _ body: () -> Component) {
   print(body())
 }
 
+func colorWithAutoClosure(_ color: @autoclosure () -> Color) -> Color {
+  return color()
+}
+
+var trueValue = true
 acceptComponentBuilder {
   "hello"
-  if true {
+  if trueValue {
     3.14159
+    colorWithAutoClosure(.red)
   }
   .red
 }
-// CHECK: array([main.Component.string("hello"), main.Component.optional(Optional(main.Component.array([main.Component.floating(3.14159)]))), main.Component.color(main.Color.red)])
+// CHECK: array([main.Component.string("hello"), main.Component.optional(Optional(main.Component.array([main.Component.floating(3.14159), main.Component.color(main.Color.red)]))), main.Component.color(main.Color.red)])
 
 // rdar://53325810