treat arrow functions as another function type when transforming

lib/6to5/transformers/arrow-functions.js

@@ -3,17 +3,9 @@ var util     = require("../util");
 var b        = require("ast-types").builders;
 
 exports.ArrowFunctionExpression = function (node) {
-  var body = node.body;
-
-  // expression body
-  if (body.type !== "BlockStatement") {
-    body = b.blockStatement([
-      b.returnStatement(body)
-    ]);
-  }
+  util.ensureBlock(node);
 
   node.expression = false;
-  node.body = body;
   node.type = "FunctionExpression";
 
   if (traverse.hasType(node, "ThisExpression")) {
@@ -34,6 +26,8 @@ exports.FunctionExpression = function (node, parent, opts, generateUid) {
   traverse(node, function (node) {
     if (node.type !== "ArrowFunctionExpression") return;
 
+    // traverse all child nodes of this arrow function and find a sole arguments
+    // identifier
     traverse(node, function (node, parent) {
       if (node.type === "Identifier" && node.name === "arguments" &&
           parent.type !== "MemberExpression") {
@@ -41,7 +35,7 @@ exports.FunctionExpression = function (node, parent, opts, generateUid) {
         id = id || b.identifier(generateUid("arguments"));
         return id;
       }
-    }, ["FunctionDeclaration", "FunctionExpression", "ArrowFunctionExpression"]);
+    }, traverse.FUNCTION_TYPES);
 
     return false;
   }, ["FunctionDeclaration", "FunctionExpression"]);

lib/6to5/transformers/block-binding.js

@@ -3,7 +3,7 @@ var util     = require("../util");
 var b        = require("ast-types").builders;
 var _        = require("lodash");
 
-var blockTypes = ["FunctionDeclaration", "FunctionExpression", "BlockStatement"];
+var blockTypes = traverse.FUNCTION_TYPES.concat(["BlockStatement"]);
 
 var isLet = function (node) {
   if (node && node.type === "VariableDeclaration" && node.kind === "let") {
@@ -31,7 +31,7 @@ exports.BlockStatement = function (node, parent, opts, generateUid) {
   if (!hasLet(node.body)) return;
 
   // ignore if we're the body of a closure already
-  if (parent.type === "FunctionExpression" || parent.type === "FunctionDeclaration") return;
+  if (_.contains(traverse.FUNCTION_TYPES, parent.type)) return;
 
   var body = node.body;
 
@@ -120,7 +120,7 @@ var buildNode = function (node) {
 
   var templateName = "function-call";
   if (traverse.hasType(node, "ThisExpression")) templateName += "-this";
-  if (traverse.hasType(node, "ReturnStatement", ["FunctionDeclaration", "FunctionExpression"])) templateName += "-return";
+  if (traverse.hasType(node, "ReturnStatement", traverse.FUNCTION_TYPES)) templateName += "-return";
 
   //
 

lib/6to5/transformers/default-parameters.js

@@ -1,8 +1,12 @@
 var util = require("../util");
 var _    = require("lodash");
 
+exports.ArrowFunctionExpression =
 exports.FunctionDeclaration =
 exports.FunctionExpression = function (node) {
+  if (!node.defaults.length) return;
+  util.ensureBlock(node);
+
   _.each(node.defaults, function (def, i) {
     if (!def) return;
 

lib/6to5/transformers/destructuring.js

@@ -85,21 +85,25 @@ exports.ForOfStatement = function (node, parent, opts, generateUid) {
   block.body = nodes.concat(block.body);
 };
 
+exports.ArrowFunctionExpression =
 exports.FunctionDeclaration =
 exports.FunctionExpression = function (node, parent, opts, generateUid) {
   var block = node.body;
   var nodes = [];
 
+  var hasDestructuring = false;
+
   node.params = node.params.map(function (pattern) {
     if (!util.isPattern(pattern)) return pattern;
 
+    hasDestructuring = true;
     var parentId = b.identifier(generateUid("ref"));
-
     pushPattern("var", nodes, pattern, parentId, generateUid);
-
     return parentId;
   });
 
+  if (!hasDestructuring) return;
+  util.ensureBlock(node);
   block.body = nodes.concat(block.body || []);
 };
 

lib/6to5/transformers/rest-parameters.js

@@ -1,6 +1,7 @@
 var util = require("../util");
 var b    = require("ast-types").builders;
 
+exports.ArrowFunctionExpression =
 exports.FunctionDeclaration =
 exports.FunctionExpression = function (node) {
   if (!node.rest) return;
@@ -11,6 +12,7 @@ exports.FunctionExpression = function (node) {
   var templateName = "arguments-slice-assign";
   if (node.params.length) templateName += "-arg";
 
+  util.ensureBlock(node);
   node.body.body.unshift(util.template(templateName, {
     VARIABLE_NAME: rest,
     SLICE_ARG: b.literal(node.params.length)

lib/6to5/traverse/index.js

@@ -64,6 +64,8 @@ var traverse = module.exports = function (parent, callbacks, blacklistTypes) {
   });
 };
 
+traverse.FUNCTION_TYPES = ["ArrowFunctionExpression", "FunctionDeclaration", "FunctionExpression"];
+
 traverse.Delete = {};
 
 traverse.hasType = function (tree, type, blacklistTypes) {

lib/6to5/util.js

@@ -1,15 +1,23 @@
 var traverse = require("./traverse");
+var astTypes = require("ast-types");
 var recast   = require("recast");
 var path     = require("path");
 var fs       = require("fs");
 var _        = require("lodash");
-var b        = require("ast-types").builders;
+
+var n = astTypes.namedTypes;
+var b = astTypes.builders;
 
 exports.ensureBlock = function (node) {
   var block = node.body;
   if (block.type === "BlockStatement") return;
 
-  node.body = [node.body];
+  if (!_.isArray(block)) {
+    if (!n.Statement.check(block)) block = b.returnStatement(block);
+    block = [block];
+  }
+
+  node.body = b.blockStatement(block);
 };
 
 exports.isPattern = function (node) {