From 5f21fc2f17f5466fcf22893f255dfd3e64e77814 Mon Sep 17 00:00:00 2001
From: Sebastian McKenzie <sebmck@gmail.com>
Date: Sun, 30 Nov 2014 11:06:31 +1100
Subject: [PATCH] ensureBlock on multiple node replacements - fixes #229

---
 lib/6to5/traverse/index.js | 4 ++++
 lib/6to5/types/index.js    | 7 +++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/lib/6to5/traverse/index.js b/lib/6to5/traverse/index.js
index 227bee3bd5..76ae72e24e 100644
--- a/lib/6to5/traverse/index.js
+++ b/lib/6to5/traverse/index.js
@@ -50,6 +50,10 @@ function traverse(parent, callbacks, opts) {
         if (result != null) {
           updated = true;
           node = obj[key] = result;
+
+          if (_.contains(t.STATEMENT_OR_BLOCK_KEYS, key) && !t.isBlockStatement(obj)) {
+            t.ensureBlock(obj, key);
+          }
         }
       };
 
diff --git a/lib/6to5/types/index.js b/lib/6to5/types/index.js
index 5db1a528d5..8a4f859e2f 100644
--- a/lib/6to5/types/index.js
+++ b/lib/6to5/types/index.js
@@ -12,6 +12,8 @@ var addAssert = function (type, is) {
   };
 };
 
+t.STATEMENT_OR_BLOCK_KEYS = ["consequent", "body"];
+
 //
 
 t.VISITOR_KEYS = require("./visitor-keys");
@@ -167,8 +169,9 @@ t.isValidIdentifier = function (name) {
   return _.isString(name) && esutils.keyword.isIdentifierName(name) && !esutils.keyword.isKeywordES6(name, true);
 };
 
-t.ensureBlock = function (node) {
-  node.body = t.toBlock(node.body, node);
+t.ensureBlock = function (node, key) {
+  key = key || "body";
+  node[key] = t.toBlock(node[key], node);
 };
 
 t.toStatement = function (node, ignore) {