Path#ensureBlock keeps path context (#6337)

* Path#ensureBlock keeps path context

This ensures that if you're inside an ArrowFunction with an expression body (say, you're on the BooleanLiteral in `() => true`), you don't suddenly lose your path context after inserting a variable.

This is because of 82d8aded8e (diff-9e0668ad44535be897b934e7077ecea5R14). Basically, an innocent `Scope#push` caused my visitor to suddenly stop working. Now, we mutate the Path so it's still in the tree.

* Tests

packages/babel-plugin-check-es2015-constants/test/fixtures/general/no-for-in/expected.js

@@ -2,4 +2,5 @@ var MULTIPLIER = 5;
 
 for (MULTIPLIER in arr) {
   throw new Error("\"MULTIPLIER\" is read-only");
+  ;
 }

packages/babel-plugin-transform-class-properties/src/index.js

@@ -193,16 +193,6 @@ export default function({ types: t }) {
 
         path.insertAfter(nodes);
       },
-      ArrowFunctionExpression(path) {
-        const classExp = path.get("body");
-        if (!classExp.isClassExpression()) return;
-
-        const body = classExp.get("body");
-        const members = body.get("body");
-        if (members.some(member => member.isClassProperty())) {
-          path.ensureBlock();
-        }
-      },
     },
   };
 }

packages/babel-plugin-transform-es2015-block-scoping/src/index.js

@@ -44,8 +44,8 @@ export default function() {
       },
 
       Loop(path, file) {
-        const { node, parent, scope } = path;
-        t.ensureBlock(node);
+        const { parent, scope } = path;
+        path.ensureBlock();
         const blockScoping = new BlockScoping(
           path,
           path.get("body"),

packages/babel-plugin-transform-es2015-parameters/test/fixtures/parameters/rest-member-expression-deoptimisation/expected.js

@@ -160,6 +160,7 @@ function forOf() {
   try {
     for (var _iterator = this[Symbol.iterator](), _step; !(_iteratorNormalCompletion = (_step = _iterator.next()).done); _iteratorNormalCompletion = true) {
       rest[0] = _step.value;
+      ;
     }
   } catch (err) {
     _didIteratorError = true;

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

@@ -23,7 +23,50 @@ export function toComputedKey(): Object {
 }
 
 export function ensureBlock() {
-  return t.ensureBlock(this.node);
+  const body = this.get("body");
+  const bodyNode = body.node;
+
+  if (Array.isArray(body)) {
+    throw new Error("Can't convert array path to a block statement");
+  }
+  if (!bodyNode) {
+    throw new Error("Can't convert node without a body");
+  }
+
+  if (body.isBlockStatement()) {
+    return bodyNode;
+  }
+
+  const statements = [];
+
+  let stringPath = "body";
+  let key;
+  let listKey;
+  if (body.isStatement()) {
+    listKey = "body";
+    key = 0;
+    statements.push(body.node);
+  } else {
+    stringPath += ".body.0";
+    if (this.isFunction()) {
+      key = "argument";
+      statements.push(t.returnStatement(body.node));
+    } else {
+      key = "expression";
+      statements.push(t.expressionStatement(body.node));
+    }
+  }
+
+  this.node.body = t.blockStatement(statements);
+  const parentPath = this.get(stringPath);
+  body.setup(
+    parentPath,
+    listKey ? parentPath.node[listKey] : parentPath.node,
+    listKey,
+    key,
+  );
+
+  return this.node;
 }
 
 /**

packages/babel-traverse/src/scope/index.js

@@ -799,7 +799,7 @@ export default class Scope {
     }
 
     if (path.isLoop() || path.isCatchClause() || path.isFunction()) {
-      t.ensureBlock(path.node);
+      path.ensureBlock();
       path = path.get("body");
     }
 

packages/babel-traverse/test/conversion.js

@@ -0,0 +1,69 @@
+import traverse from "../lib";
+import assert from "assert";
+import { parse } from "babylon";
+import generate from "babel-generator";
+import * as t from "babel-types";
+
+function getPath(code) {
+  const ast = parse(code);
+  let path;
+  traverse(ast, {
+    Program: function(_path) {
+      path = _path.get("body.0");
+      _path.stop();
+    },
+  });
+
+  return path;
+}
+
+function generateCode(path) {
+  return generate(path.parentPath.node).code;
+}
+
+describe("conversion", function() {
+  describe("ensureBlock", function() {
+    it("throws converting node without body to block", function() {
+      const rootPath = getPath("true;");
+
+      assert.throws(() => {
+        rootPath.ensureBlock();
+      }, /Can't convert node without a body/);
+    });
+
+    it("throws converting already block array", function() {
+      const rootPath = getPath("function test() { true; }").get("body");
+      assert.throws(() => {
+        rootPath.ensureBlock();
+      }, /Can't convert array path to a block statement/);
+    });
+
+    it("converts arrow function with expression body to block", function() {
+      const rootPath = getPath("() => true").get("expression");
+      rootPath.ensureBlock();
+      assert.equal(generateCode(rootPath), "() => {\n  return true;\n};");
+    });
+
+    it("preserves arrow function body's context", function() {
+      const rootPath = getPath("() => true").get("expression");
+      const body = rootPath.get("body");
+      rootPath.ensureBlock();
+      body.replaceWith(t.booleanLiteral(false));
+      assert.equal(generateCode(rootPath), "() => {\n  return false;\n};");
+    });
+
+    it("converts for loop with statement body to block", function() {
+      const rootPath = getPath("for (;;) true;");
+      rootPath.ensureBlock();
+      assert.equal(generateCode(rootPath), "for (;;) {\n  true;\n}");
+    });
+
+    it("preserves for loop body's context", function() {
+      const rootPath = getPath("for (;;) true;");
+      const body = rootPath.get("body");
+      rootPath.ensureBlock();
+      body.replaceWith(t.booleanLiteral(false));
+      assert.equal(generateCode(rootPath), "for (;;) {\n  false;\n}");
+    });
+  });
+});