Optimize removal-hooks for ArrowFunctions (babel#5076)

packages/babel-traverse/package.json

@@ -17,5 +17,8 @@
     "globals": "^9.0.0",
     "invariant": "^2.2.0",
     "lodash": "^4.2.0"
+  },
+  "devDependencies": {
+    "babel-generator": "^6.21.0"
   }
 }

packages/babel-traverse/src/path/lib/removal-hooks.js

@@ -5,13 +5,6 @@
  */
 
 export let hooks = [
-  function (self, parent) {
-    if (self.key === "body" && parent.isArrowFunctionExpression()) {
-      self.replaceWith(self.scope.buildUndefinedNode());
-      return true;
-    }
-  },
-
   function (self, parent) {
     let removeParent = false;
 
@@ -69,7 +62,7 @@ export let hooks = [
   function (self, parent) {
     if (
       (parent.isIfStatement() && (self.key === "consequent" || self.key === "alternate")) ||
-      (parent.isLoop() && self.key === "body")
+      (self.key === "body" && (parent.isLoop() || parent.isArrowFunctionExpression()))
     ) {
       self.replaceWith({
         type: "BlockStatement",

packages/babel-traverse/test/removal.js

@@ -0,0 +1,34 @@
+import traverse from "../lib";
+import assert from "assert";
+import { parse } from "babylon";
+import generate from "babel-generator";
+
+function getPath(code) {
+  const ast = parse(code);
+  let path;
+  traverse(ast, {
+    Program: function (_path) {
+      path = _path;
+      _path.stop();
+    }
+  });
+
+  return path;
+}
+
+function generateCode(path) {
+  return generate(path.node).code;
+}
+
+describe("removal", function () {
+  describe("ArrowFunction", function () {
+    it("remove body", function () {
+      const rootPath = getPath("x = () => b;");
+      const path = rootPath.get("body")[0].get("expression").get("right");
+      const body = path.get("body");
+      body.remove();
+
+      assert.equal(generateCode(rootPath), "x = () => {};", "body should be replaced with BlockStatement");
+    });
+  });
+});