Merge pull request swiftlang#32100 from DougGregor/function-builders-do-block

[Function builders] Align buildDo() implementation with the pitch.

Author: DougGregor

include/swift/AST/Stmt.h

@@ -539,11 +539,11 @@ class LabeledStmt : public Stmt {
 /// DoStmt - do statement, without any trailing clauses.
 class DoStmt : public LabeledStmt {
   SourceLoc DoLoc;
-  Stmt *Body;
+  BraceStmt *Body;
 
 public:
   DoStmt(LabeledStmtInfo labelInfo, SourceLoc doLoc,
-         Stmt *body, Optional<bool> implicit = None)
+         BraceStmt *body, Optional<bool> implicit = None)
     : LabeledStmt(StmtKind::Do, getDefaultImplicitFlag(implicit, doLoc),
                   labelInfo),
       DoLoc(doLoc), Body(body) {}
@@ -553,8 +553,8 @@ class DoStmt : public LabeledStmt {
   SourceLoc getStartLoc() const { return getLabelLocOrKeywordLoc(DoLoc); }
   SourceLoc getEndLoc() const { return Body->getEndLoc(); }
 
-  Stmt *getBody() const { return Body; }
-  void setBody(Stmt *s) { Body = s; }
+  BraceStmt *getBody() const { return Body; }
+  void setBody(BraceStmt *s) { Body = s; }
 
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Do; }
 };

lib/AST/ASTWalker.cpp

@@ -1493,7 +1493,7 @@ Stmt *Traversal::visitGuardStmt(GuardStmt *US) {
 }
 
 Stmt *Traversal::visitDoStmt(DoStmt *DS) {
-  if (Stmt *S2 = doIt(DS->getBody()))
+  if (BraceStmt *S2 = cast_or_null<BraceStmt>(doIt(DS->getBody())))
     DS->setBody(S2);
   else
     return nullptr;

lib/Sema/BuilderTransform.cpp

@@ -293,6 +293,10 @@ class BuilderClosureVisitor
   }
 
   VarDecl *visitBraceStmt(BraceStmt *braceStmt) {
+    return visitBraceStmt(braceStmt, ctx.Id_buildBlock);
+  }
+
+  VarDecl *visitBraceStmt(BraceStmt *braceStmt, Identifier builderFunction) {
     SmallVector<Expr *, 4> expressions;
     auto addChild = [&](VarDecl *childVar) {
       if (!childVar)
@@ -359,7 +363,7 @@ class BuilderClosureVisitor
 
     // Call Builder.buildBlock(... args ...)
     auto call = buildCallIfWanted(braceStmt->getStartLoc(),
-                                  ctx.Id_buildBlock, expressions,
+                                  builderFunction, expressions,
                                   /*argLabels=*/{ });
     if (!call)
       return nullptr;
@@ -380,17 +384,13 @@ class BuilderClosureVisitor
       return nullptr;
     }
 
-    auto childVar = visit(doStmt->getBody());
+    auto childVar = visitBraceStmt(doStmt->getBody(), ctx.Id_buildDo);
     if (!childVar)
       return nullptr;
 
     auto childRef = buildVarRef(childVar, doStmt->getEndLoc());
-    auto call = buildCallIfWanted(doStmt->getStartLoc(), ctx.Id_buildDo,
-                                  childRef, /*argLabels=*/{ });
-    if (!call)
-      return nullptr;
 
-    return captureExpr(call, /*oneWay=*/true, doStmt);
+    return captureExpr(childRef, /*oneWay=*/true, doStmt);
   }
 
   CONTROL_FLOW_STMT(Yield)

lib/Sema/TypeCheckStmt.cpp

@@ -658,7 +658,7 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
 
   Stmt *visitDoStmt(DoStmt *DS) {
     AddLabeledStmt loopNest(*this, DS);
-    Stmt *S = DS->getBody();
+    BraceStmt *S = DS->getBody();
     typeCheckStmt(S);
     DS->setBody(S);
     return DS;

test/Constraints/function_builder.swift

@@ -6,6 +6,10 @@ enum Either<T,U> {
   case second(U)
 }
 
+struct Do<T> {
+  var value: T
+}
+
 @_functionBuilder
 struct TupleBuilder {
   static func buildBlock<T1>(_ t1: T1) -> (T1) {
@@ -32,7 +36,19 @@ struct TupleBuilder {
     return (t1, t2, t3, t4, t5)
   }
 
-  static func buildDo<T>(_ value: T) -> T { return value }
+  static func buildDo<T1>(_ t1: T1) -> Do<(T1)> {
+    .init(value: t1)
+  }
+
+  static func buildDo<T1, T2>(_ t1: T1, _ t2: T2) -> Do<(T1, T2)> {
+    .init(value: (t1, t2))
+  }
+  
+  static func buildDo<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
+      -> Do<(T1, T2, T3)> {
+    .init(value: (t1, t2, t3))
+  }
+
   static func buildIf<T>(_ value: T?) -> T? { return value }
 
   static func buildEither<T,U>(first value: T) -> Either<T,U> {
@@ -49,7 +65,7 @@ func tuplify<T>(_ cond: Bool, @TupleBuilder body: (Bool) -> T) {
   print(body(cond))
 }
 
-// CHECK: (17, 3.14159, "Hello, DSL", (["nested", "do"], 6), Optional((2.71828, ["if", "stmt"])))
+// CHECK: (17, 3.14159, "Hello, DSL", main.Do<(Swift.Array<Swift.String>, Swift.Int)>(value: (["nested", "do"], 6)), Optional((2.71828, ["if", "stmt"])))
 let name = "dsl"
 tuplify(true) {
   17

test/ModuleInterface/function_builders.swift

@@ -26,7 +26,15 @@ public struct TupleBuilder {
     return (t1, t2, t3, t4, t5)
   }
 
-  public static func buildDo<T>(_ value: T) -> T { return value }
+  public static func buildDo<T1, T2>(_ t1: T1, _ t2: T2) -> (T1, T2) {
+    return (t1, t2)
+  }
+
+  public static func buildDo<T1, T2, T3>(_ t1: T1, _ t2: T2, _ t3: T3)
+      -> (T1, T2, T3) {
+    return (t1, t2, t3)
+  }
+
   public static func buildIf<T>(_ value: T?) -> T? { return value }
 }