Replace arrow expression body with block statement

Original PR: https://github.com/babel/babel/pull/2469. Seems this got
lost in the v6 changes.

- - -

Without this, the only way to replace the arrow function is to either
manually override its `node.body`, or duplicate the arrow:

```js
// Old
ArrowFunctionExpression: function (node) {
  node.body = t.blockStatement(...);
  // Or
  return t.ArrowFunctionExpression(
    node.params,
    t.blockStatement(...),
    node.async
  );
}

// New
ArrowFunctionExpression: function() {
  this.get("body").replaceWith(t.blockStatement(...));
}
```

packages/babel-core/test/path.js

@@ -19,4 +19,138 @@ suite("traversal path", function () {
 
     chai.expect(actualCode).to.be.equal("console.whatever();");
   });
+
+  test("replaceWith (arrow expression body to block statement body)", function () {
+    var expectCode = "var fn = () => true;";
+
+    var actualCode = transform(expectCode, {
+      plugins: [new Plugin({
+        visitor: {
+          ArrowFunctionExpression: function (path) {
+            path.get("body").replaceWith({
+              type: "BlockStatement",
+              body: [{
+                type: "ReturnStatement",
+                argument: {
+                  type: "BooleanLiteral",
+                  value: true
+                }
+              }]
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("var fn = () => {\n  return true;\n};");
+  });
+
+  test("replaceWith (arrow block statement body to expression body)", function () {
+    var expectCode = "var fn = () => { return true; }";
+
+    var actualCode = transform(expectCode, {
+      plugins: [new Plugin({
+        visitor: {
+          ArrowFunctionExpression: function (path) {
+            path.get("body").replaceWith({
+              type: "BooleanLiteral",
+              value: true
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("var fn = () => true;");
+  });
+
+  test("replaceWith (for-in left expression to variable declaration)", function () {
+    var expectCode = "for (KEY in right);";
+
+    var actualCode = transform(expectCode, {
+      plugins: [new Plugin({
+        visitor: {
+          ForInStatement: function (path) {
+            path.get("left").replaceWith({
+              type: "VariableDeclaration",
+              kind: "var",
+              declarations: [{
+                type: "VariableDeclarator",
+                id: {
+                  type: "Identifier",
+                  name: "KEY"
+                }
+              }]
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("for (var KEY in right);");
+  });
+
+  test("replaceWith (for-in left variable declaration to expression)", function () {
+    var expectCode = "for (var KEY in right);";
+
+    var actualCode = transform(expectCode, {
+      plugins: [new Plugin({
+        visitor: {
+          ForInStatement: function (path) {
+            path.get("left").replaceWith({
+              type: "Identifier",
+              name: "KEY"
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("for (KEY in right);");
+  });
+
+  test("replaceWith (for-loop left expression to variable declaration)", function () {
+    var expectCode = "for (KEY;;);";
+
+    var actualCode = transform(expectCode, {
+      plugins: [new Plugin({
+        visitor: {
+          ForStatement: function (path) {
+            path.get("init").replaceWith({
+              type: "VariableDeclaration",
+              kind: "var",
+              declarations: [{
+                type: "VariableDeclarator",
+                id: {
+                  type: "Identifier",
+                  name: "KEY"
+                }
+              }]
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("for (var KEY;;);");
+  });
+
+  test("replaceWith (for-loop left variable declaration to expression)", function () {
+    var expectCode = "for (var KEY;;);";
+
+    var actualCode = transform(expectCode, {
+      plugins: [new Plugin({
+        visitor: {
+          ForStatement: function (path) {
+            path.get("init").replaceWith({
+              type: "Identifier",
+              name: "KEY"
+            });
+          }
+        }
+      })]
+    }).code;
+
+    chai.expect(actualCode).to.be.equal("for (KEY;;);");
+  });
 });

packages/babel-traverse/src/path/introspection.js

@@ -116,7 +116,7 @@ export function isNodeType(type: string): boolean {
 }
 
 /**
- * This checks whether or now we're in one of the following positions:
+ * This checks whether or not we're in one of the following positions:
  *
  *   for (KEY in right);
  *   for (KEY;;);
@@ -129,6 +129,28 @@ export function canHaveVariableDeclarationOrExpression() {
   return (this.key === "init" || this.key === "left") && this.parentPath.isFor();
 }
 
+/**
+ * This checks whether we are swapping an arrow function's body between an
+ * expression and a block statement (or vice versa).
+ *
+ * This is because arrow functions may implicitly return an expression, which
+ * is the same as containing a block statement.
+ */
+
+export function canSwapBetweenExpressionAndStatement(replacement) {
+  if (this.key !== "body" || !this.parentPath.isArrowFunctionExpression()) {
+    return false;
+  }
+
+  if (this.isExpression()) {
+    return t.isBlockStatement(replacement);
+  } else if (this.isBlockStatement()) {
+    return t.isExpression(replacement);
+  }
+
+  return false;
+}
+
 /**
  * Check whether the current path references a completion record
  */

packages/babel-traverse/src/path/replacement.js

@@ -121,14 +121,18 @@ export function replaceWith(replacement) {
     throw new Error("Don't use `path.replaceWith()` with a source string, use `path.replaceWithSourceString()`");
   }
 
-  // replacing a statement with an expression so wrap it in an expression statement
-  if (this.isNodeType("Statement") && t.isExpression(replacement) && !this.canHaveVariableDeclarationOrExpression()) {
-    replacement = t.expressionStatement(replacement);
+  if (this.isNodeType("Statement") && t.isExpression(replacement)) {
+    if (!this.canHaveVariableDeclarationOrExpression() && !this.canSwapBetweenExpressionAndStatement(replacement)) {
+      // replacing a statement with an expression so wrap it in an expression statement
+      replacement = t.expressionStatement(replacement);
+    }
   }
 
-  // replacing an expression with a statement so let's explode it
   if (this.isNodeType("Expression") && t.isStatement(replacement)) {
-    return this.replaceExpressionWithStatements([replacement]);
+    if (!this.canHaveVariableDeclarationOrExpression() && !this.canSwapBetweenExpressionAndStatement(replacement)) {
+      // replacing an expression with a statement so let's explode it
+      return this.replaceExpressionWithStatements([replacement]);
+    }
   }
 
   let oldNode = this.node;