[CS] Factor out brace element handling in a couple of cases

Split out brace element handling for a couple of
walkers into a new function. The diff for this is
best viewed without whitespace changes. This
should be NFC.

Author: hamishknight

lib/Sema/BuilderTransform.cpp

@@ -364,68 +364,70 @@ class BuilderClosureVisitor
     }
   }
 
-  VarDecl *visitBraceStmt(BraceStmt *braceStmt) {
-    SmallVector<Expr *, 4> expressions;
+  void visitBraceElement(ASTNode node, SmallVectorImpl<Expr *> &expressions) {
     auto addChild = [&](VarDecl *childVar) {
       if (!childVar)
         return;
 
       expressions.push_back(builder.buildVarRef(childVar, childVar->getLoc()));
     };
 
-    for (auto node : braceStmt->getElements()) {
-      // Implicit returns in single-expression function bodies are treated
-      // as the expression.
-      if (auto returnStmt =
-              dyn_cast_or_null<ReturnStmt>(node.dyn_cast<Stmt *>())) {
-        assert(returnStmt->isImplicit());
-        node = returnStmt->getResult();
-      }
-
-      if (auto stmt = node.dyn_cast<Stmt *>()) {
-        addChild(visit(stmt));
-        continue;
-      }
+    // Implicit returns in single-expression function bodies are treated
+    // as the expression.
+    if (auto returnStmt =
+            dyn_cast_or_null<ReturnStmt>(node.dyn_cast<Stmt *>())) {
+      assert(returnStmt->isImplicit());
+      node = returnStmt->getResult();
+    }
 
-      if (auto decl = node.dyn_cast<Decl *>()) {
-        // Just ignore #if; the chosen children should appear in the
-        // surrounding context.  This isn't good for source tools but it
-        // at least works.
-        if (isa<IfConfigDecl>(decl))
-          continue;
+    if (auto stmt = node.dyn_cast<Stmt *>()) {
+      addChild(visit(stmt));
+      return;
+    }
 
-        // Skip #warning/#error; we'll handle them when applying the builder.
-        if (isa<PoundDiagnosticDecl>(decl)) {
-          continue;
-        }
+    if (auto decl = node.dyn_cast<Decl *>()) {
+      // Just ignore #if; the chosen children should appear in the
+      // surrounding context.  This isn't good for source tools but it
+      // at least works.
+      if (isa<IfConfigDecl>(decl))
+        return;
 
-        // Pattern bindings are okay so long as all of the entries are
-        // initialized.
-        if (auto patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
-          visitPatternBindingDecl(patternBinding);
-          continue;
-        }
+      // Skip #warning/#error; we'll handle them when applying the builder.
+      if (isa<PoundDiagnosticDecl>(decl))
+        return;
 
-        // Ignore variable declarations, because they're always handled within
-        // their enclosing pattern bindings.
-        if (isa<VarDecl>(decl))
-          continue;
+      // Pattern bindings are okay so long as all of the entries are
+      // initialized.
+      if (auto patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
+        visitPatternBindingDecl(patternBinding);
+        return;
+      }
 
-        if (!unhandledNode)
-          unhandledNode = decl;
+      // Ignore variable declarations, because they're always handled within
+      // their enclosing pattern bindings.
+      if (isa<VarDecl>(decl))
+        return;
 
-        continue;
-      }
+      if (!unhandledNode)
+        unhandledNode = decl;
 
-      auto expr = node.get<Expr *>();
-      if (cs && builder.supports(ctx.Id_buildExpression)) {
-        expr = buildCallIfWanted(expr->getLoc(), ctx.Id_buildExpression,
-                                 { expr }, { Identifier() });
-      }
+      return;
+    }
 
-      addChild(captureExpr(expr, /*oneWay=*/true, node.get<Expr *>()));
+    auto expr = node.get<Expr *>();
+    if (cs && builder.supports(ctx.Id_buildExpression)) {
+      expr = buildCallIfWanted(expr->getLoc(), ctx.Id_buildExpression,
+                               {expr}, {Identifier()});
     }
 
+    addChild(captureExpr(expr, /*oneWay=*/true, node.get<Expr *>()));
+  }
+
+  VarDecl *visitBraceStmt(BraceStmt *braceStmt) {
+    SmallVector<Expr *, 4> expressions;
+    for (auto node : braceStmt->getElements())
+      visitBraceElement(node, expressions);
+
     if (!cs || hadError)
       return nullptr;
 
@@ -1769,107 +1771,113 @@ class BuilderClosureRewriter
         solution(solution), dc(dc), builderTransform(builderTransform),
         rewriteTarget(rewriteTarget) { }
 
-  NullablePtr<Stmt>
-  visitBraceStmt(BraceStmt *braceStmt, ResultBuilderTarget target,
-                 Optional<ResultBuilderTarget> innerTarget = None) {
-    std::vector<ASTNode> newElements;
-
-    // If there is an "inner" target corresponding to this brace, declare
-    // it's temporary variable if needed.
-    if (innerTarget) {
-      declareTemporaryVariable(innerTarget->captured.first, newElements);
+  /// Visit the element of a brace statement, returning \c false if the element
+  /// was rewritten successfully, or \c true if there was an error.
+  bool visitBraceElement(ASTNode node, std::vector<ASTNode> &newElements) {
+    // Implicit returns in single-expression function bodies are treated
+    // as the expression.
+    if (auto returnStmt =
+            dyn_cast_or_null<ReturnStmt>(node.dyn_cast<Stmt *>())) {
+      assert(returnStmt->isImplicit());
+      node = returnStmt->getResult();
     }
 
-    for (auto node : braceStmt->getElements()) {
-      // Implicit returns in single-expression function bodies are treated
-      // as the expression.
-      if (auto returnStmt =
-              dyn_cast_or_null<ReturnStmt>(node.dyn_cast<Stmt *>())) {
-        assert(returnStmt->isImplicit());
-        node = returnStmt->getResult();
-      }
+    if (auto expr = node.dyn_cast<Expr *>()) {
+      // Skip error expressions.
+      if (isa<ErrorExpr>(expr))
+        return false;
 
-      if (auto expr = node.dyn_cast<Expr *>()) {
-        // Skip error expressions.
-        if (isa<ErrorExpr>(expr))
-          continue;
+      // Each expression turns into a 'let' that captures the value of
+      // the expression.
+      auto recorded = takeCapturedExpr(expr);
 
-        // Each expression turns into a 'let' that captures the value of
-        // the expression.
-        auto recorded = takeCapturedExpr(expr);
+      // Rewrite the expression
+      Expr *finalExpr = rewriteExpr(recorded.generatedExpr);
 
-        // Rewrite the expression
-        Expr *finalExpr = rewriteExpr(recorded.generatedExpr);
+      // Form a new pattern binding to bind the temporary variable to the
+      // transformed expression.
+      declareTemporaryVariable(recorded.temporaryVar, newElements, finalExpr);
+      return false;
+    }
 
-        // Form a new pattern binding to bind the temporary variable to the
-        // transformed expression.
-        declareTemporaryVariable(recorded.temporaryVar, newElements, finalExpr);
-        continue;
+    if (auto stmt = node.dyn_cast<Stmt *>()) {
+      // "throw" statements produce no value. Transform them directly.
+      if (auto throwStmt = dyn_cast<ThrowStmt>(stmt)) {
+        if (auto newStmt = visitThrowStmt(throwStmt)) {
+          newElements.push_back(newStmt.get());
+        }
+        return false;
       }
 
-      if (auto stmt = node.dyn_cast<Stmt *>()) {
-        // "throw" statements produce no value. Transform them directly.
-        if (auto throwStmt = dyn_cast<ThrowStmt>(stmt)) {
-          if (auto newStmt = visitThrowStmt(throwStmt)) {
-            newElements.push_back(newStmt.get());
-          }
-          continue;
-        }
+      // Each statement turns into a (potential) temporary variable
+      // binding followed by the statement itself.
+      auto captured = takeCapturedStmt(stmt);
 
-        // Each statement turns into a (potential) temporary variable
-        // binding followed by the statement itself.
-        auto captured = takeCapturedStmt(stmt);
+      declareTemporaryVariable(captured.first, newElements);
 
-        declareTemporaryVariable(captured.first, newElements);
+      auto finalStmt =
+          visit(stmt, ResultBuilderTarget{ResultBuilderTarget::TemporaryVar,
+                                          std::move(captured)});
 
-        auto finalStmt = visit(
-            stmt,
-            ResultBuilderTarget{ResultBuilderTarget::TemporaryVar,
-                                  std::move(captured)});
+      // Re-write of statements that envolve type-checking
+      // could fail, such a failure terminates the walk.
+      if (!finalStmt)
+        return true;
 
-        // Re-write of statements that envolve type-checking
-        // could fail, such a failure terminates the walk.
-        if (!finalStmt)
-          return nullptr;
+      newElements.push_back(finalStmt.get());
+      return false;
+    }
 
-        newElements.push_back(finalStmt.get());
-        continue;
-      }
+    auto decl = node.get<Decl *>();
 
-      auto decl = node.get<Decl *>();
+    // Skip #if declarations.
+    if (isa<IfConfigDecl>(decl)) {
+      newElements.push_back(decl);
+      return false;
+    }
 
-      // Skip #if declarations.
-      if (isa<IfConfigDecl>(decl)) {
-        newElements.push_back(decl);
-        continue;
-      }
+    // Diagnose #warning / #error during application.
+    if (auto poundDiag = dyn_cast<PoundDiagnosticDecl>(decl)) {
+      TypeChecker::typeCheckDecl(poundDiag);
+      newElements.push_back(decl);
+      return false;
+    }
 
-      // Diagnose #warning / #error during application.
-      if (auto poundDiag = dyn_cast<PoundDiagnosticDecl>(decl)) {
-        TypeChecker::typeCheckDecl(poundDiag);
-        newElements.push_back(decl);
-        continue;
-      }
+    // Skip variable declarations; they're always part of a pattern
+    // binding.
+    if (isa<VarDecl>(decl)) {
+      TypeChecker::typeCheckDecl(decl);
+      newElements.push_back(decl);
+      return false;
+    }
 
-      // Skip variable declarations; they're always part of a pattern
-      // binding.
-      if (isa<VarDecl>(decl)) {
-        TypeChecker::typeCheckDecl(decl);
-        newElements.push_back(decl);
-        continue;
-      }
+    // Handle pattern bindings.
+    if (auto patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
+      auto resultTarget =
+          rewriteTarget(SolutionApplicationTarget{patternBinding});
+      assert(resultTarget.has_value() &&
+             "Could not rewrite pattern binding entries!");
+      TypeChecker::typeCheckDecl(resultTarget->getAsPatternBinding());
+      newElements.push_back(resultTarget->getAsPatternBinding());
+      return false;
+    }
 
-      // Handle pattern bindings.
-      if (auto patternBinding = dyn_cast<PatternBindingDecl>(decl)) {
-        auto resultTarget = rewriteTarget(SolutionApplicationTarget{patternBinding});
-        assert(resultTarget.has_value()
-               && "Could not rewrite pattern binding entries!");
-        TypeChecker::typeCheckDecl(resultTarget->getAsPatternBinding());
-        newElements.push_back(resultTarget->getAsPatternBinding());
-        continue;
-      }
+    llvm_unreachable("Cannot yet handle declarations");
+  }
 
-      llvm_unreachable("Cannot yet handle declarations");
+  NullablePtr<Stmt>
+  visitBraceStmt(BraceStmt *braceStmt, ResultBuilderTarget target,
+                 Optional<ResultBuilderTarget> innerTarget = None) {
+    std::vector<ASTNode> newElements;
+
+    // If there is an "inner" target corresponding to this brace, declare
+    // it's temporary variable if needed.
+    if (innerTarget)
+      declareTemporaryVariable(innerTarget->captured.first, newElements);
+
+    for (auto node : braceStmt->getElements()) {
+      if (visitBraceElement(node, newElements))
+        return nullptr;
     }
 
     // If there is an "inner" target corresponding to this brace, initialize