[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

@@ -1000,73 +1000,85 @@ class ResultBuilderTransform
     return failTransform(stmt);                                                \
   }
 
-  std::pair<NullablePtr<VarDecl>, Optional<UnsupportedElt>>
-  transform(BraceStmt *braceStmt, SmallVectorImpl<ASTNode> &newBody) {
-    SmallVector<Expr *, 4> buildBlockArguments;
-
-    auto failTransform = [&](UnsupportedElt unsupported) {
-      return std::make_pair(nullptr, unsupported);
-    };
-
-    for (auto element : braceStmt->getElements()) {
-      if (auto *returnStmt = getAsStmt<ReturnStmt>(element)) {
-        assert(returnStmt->isImplicit());
-        element = returnStmt->getResult();
-      }
-
-      if (auto *decl = element.dyn_cast<Decl *>()) {
-        switch (decl->getKind()) {
+  /// Visit the element of a brace statement, returning \c None if the element
+  /// was transformed successfully, or an unsupported element if the element
+  /// cannot be handled.
+  Optional<UnsupportedElt>
+  transformBraceElement(ASTNode element, SmallVectorImpl<ASTNode> &newBody,
+                        SmallVectorImpl<Expr *> &buildBlockArguments) {
+    if (auto *returnStmt = getAsStmt<ReturnStmt>(element)) {
+      assert(returnStmt->isImplicit());
+      element = returnStmt->getResult();
+    }
+
+    if (auto *decl = element.dyn_cast<Decl *>()) {
+      switch (decl->getKind()) {
         // Just ignore #if; the chosen children should appear in
         // the surrounding context.  This isn't good for source
         // tools but it at least works.
-        case DeclKind::IfConfig:
+      case DeclKind::IfConfig:
         // Skip #warning/#error; we'll handle them when applying
         // the builder.
-        case DeclKind::PoundDiagnostic:
-        case DeclKind::PatternBinding:
-        case DeclKind::Var:
-        case DeclKind::Param:
-          newBody.push_back(element);
-          break;
-
-        default:
-          return failTransform(decl);
-        }
+      case DeclKind::PoundDiagnostic:
+      case DeclKind::PatternBinding:
+      case DeclKind::Var:
+      case DeclKind::Param:
+        newBody.push_back(element);
+        return None;
+
+      default:
+        return UnsupportedElt(decl);
+      }
+      llvm_unreachable("Unhandled case in switch!");
+    }
 
-        continue;
+    if (auto *stmt = element.dyn_cast<Stmt *>()) {
+      // Throw is allowed as is.
+      if (auto *throwStmt = dyn_cast<ThrowStmt>(stmt)) {
+        newBody.push_back(throwStmt);
+        return None;
       }
 
-      if (auto *stmt = element.dyn_cast<Stmt *>()) {
-        // Throw is allowed as is.
-        if (auto *throwStmt = dyn_cast<ThrowStmt>(stmt)) {
-          newBody.push_back(throwStmt);
-          continue;
-        }
+      // Allocate variable with a placeholder type
+      auto *resultVar = buildPlaceholderVar(stmt->getStartLoc(), newBody);
 
-        // Allocate variable with a placeholder type
-        auto *resultVar = buildPlaceholderVar(stmt->getStartLoc(), newBody);
+      auto result = visit(stmt, resultVar);
+      if (!result)
+        return UnsupportedElt(stmt);
 
-        auto result = visit(stmt, resultVar);
-        if (!result)
-          return failTransform(stmt);
+      newBody.push_back(result.get());
+      buildBlockArguments.push_back(
+          builder.buildVarRef(resultVar, stmt->getStartLoc()));
+      return None;
+    }
 
-        newBody.push_back(result.get());
-        buildBlockArguments.push_back(
-            builder.buildVarRef(resultVar, stmt->getStartLoc()));
-        continue;
-      }
+    auto *expr = element.get<Expr *>();
+    if (builder.supports(ctx.Id_buildExpression)) {
+      expr = builder.buildCall(expr->getLoc(), ctx.Id_buildExpression, {expr},
+                               {Identifier()});
+    }
 
-      auto *expr = element.get<Expr *>();
-      if (builder.supports(ctx.Id_buildExpression)) {
-        expr = builder.buildCall(expr->getLoc(), ctx.Id_buildExpression, {expr},
-                                 {Identifier()});
-      }
+    auto *capture = captureExpr(expr, newBody);
+    // A reference to the synthesized variable is passed as an argument
+    // to buildBlock.
+    buildBlockArguments.push_back(
+        builder.buildVarRef(capture, element.getStartLoc()));
+    return None;
+  }
 
-      auto *capture = captureExpr(expr, newBody);
-      // A reference to the synthesized variable is passed as an argument
-      // to buildBlock.
-      buildBlockArguments.push_back(
-          builder.buildVarRef(capture, element.getStartLoc()));
+  std::pair<NullablePtr<VarDecl>, Optional<UnsupportedElt>>
+  transform(BraceStmt *braceStmt, SmallVectorImpl<ASTNode> &newBody) {
+    SmallVector<Expr *, 4> buildBlockArguments;
+
+    auto failTransform = [&](UnsupportedElt unsupported) {
+      return std::make_pair(nullptr, unsupported);
+    };
+
+    for (auto element : braceStmt->getElements()) {
+      if (auto unsupported =
+              transformBraceElement(element, newBody, buildBlockArguments)) {
+        return failTransform(*unsupported);
+      }
     }
 
     // Synthesize `buildBlock` or `buildPartial` based on captured arguments.

lib/Sema/CSSyntacticElement.cpp

@@ -1512,6 +1512,27 @@ class SyntacticElementSolutionApplication
     return caseStmt;
   }
 
+  ASTNode visitBraceElement(ASTNode node) {
+    auto &cs = solution.getConstraintSystem();
+    if (auto *expr = node.dyn_cast<Expr *>()) {
+      // Rewrite the expression.
+      auto target = *cs.getSolutionApplicationTarget(expr);
+      if (auto rewrittenTarget = rewriteTarget(target)) {
+        node = rewrittenTarget->getAsExpr();
+
+        if (target.isDiscardedExpr())
+          TypeChecker::checkIgnoredExpr(castToExpr(node));
+      } else {
+        hadError = true;
+      }
+    } else if (auto stmt = node.dyn_cast<Stmt *>()) {
+      node = visit(stmt);
+    } else {
+      visitDecl(node.get<Decl *>());
+    }
+    return node;
+  }
+
   ASTNode visitBraceStmt(BraceStmt *braceStmt) {
     auto &cs = solution.getConstraintSystem();
 
@@ -1527,24 +1548,8 @@ class SyntacticElementSolutionApplication
       }
     }
 
-    for (auto &node : braceStmt->getElements()) {
-      if (auto expr = node.dyn_cast<Expr *>()) {
-        // Rewrite the expression.
-        auto target = *cs.getSolutionApplicationTarget(expr);
-        if (auto rewrittenTarget = rewriteTarget(target)) {
-          node = rewrittenTarget->getAsExpr();
-
-          if (target.isDiscardedExpr())
-            TypeChecker::checkIgnoredExpr(castToExpr(node));
-        } else {
-          hadError = true;
-        }
-      } else if (auto stmt = node.dyn_cast<Stmt *>()) {
-        node = visit(stmt);
-      } else {
-        visitDecl(node.get<Decl *>());
-      }
-    }
+    for (auto &node : braceStmt->getElements())
+      node = visitBraceElement(node);
 
     // Source compatibility workaround.
     //