Return inserted/replaced paths (#5710)

* Return inserted/replaced paths

This gives `Path`’s replacement and insertion methods a consistent
return value: the inserted/replaced paths.

Before, they could return `undefined`, a `node`, or a the current path
inside an array. It was kinda pointless.  But now they always return an
array of paths, which is useful for solving
https://github.com/babel/babel/pull/4935#discussion_r96151368.

* Return inserted nodes and not BlockStatement

Addded test for bug #4363

* Cleanups

- `#replaceWith` will now return the current path if it's the same node
- `#insertAfter` and `#insertBefore` use public Path APIs now
- Makes container insertion faster (single splice call)
- Use public APIs in container insertion
- Replacing a statement with an expression returns the expression's path
- Replacing an expression with multiple statements returns the inserted
  closure's body's paths.

packages/babel-plugin-transform-es2015-block-scoping/test/fixtures/general/issue-4363/actual.js

@@ -0,0 +1,18 @@
+function WithoutCurlyBraces() {
+  if (true)
+    for (let k in kv) {
+        function foo() { return this }
+        function bar() { return foo.call(this) }
+        console.log(this, k) // => undefined
+    }
+}
+
+function WithCurlyBraces() {
+  if (true) {
+    for (let k in kv) {
+        function foo() { return this }
+        function bar() { return foo.call(this) }
+        console.log(this, k) // => 777
+    }
+  }
+}

packages/babel-plugin-transform-es2015-block-scoping/test/fixtures/general/issue-4363/expected.js

@@ -0,0 +1,43 @@
+function WithoutCurlyBraces() {
+  var _this = this;
+
+  if (true) {
+    var _loop = function (k) {
+      function foo() {
+        return this;
+      }
+
+      function bar() {
+        return foo.call(this);
+      }
+
+      console.log(_this, k); // => undefined
+    };
+
+    for (var k in kv) {
+      _loop(k);
+    }
+  }
+}
+
+function WithCurlyBraces() {
+  var _this2 = this;
+
+  if (true) {
+    var _loop2 = function (k) {
+      function foo() {
+        return this;
+      }
+
+      function bar() {
+        return foo.call(this);
+      }
+
+      console.log(_this2, k); // => 777
+    };
+
+    for (var k in kv) {
+      _loop2(k);
+    }
+  }
+}

packages/babel-plugin-transform-es2015-classes/src/index.js

@@ -44,7 +44,10 @@ export default function({ types: t }) {
         if (node[VISITED]) return;
 
         const inferred = nameFunction(path);
-        if (inferred && inferred !== node) return path.replaceWith(inferred);
+        if (inferred && inferred !== node) {
+          path.replaceWith(inferred);
+          return;
+        }
 
         node[VISITED] = true;
 

packages/babel-plugin-transform-es2015-for-of/src/index.js

@@ -109,7 +109,8 @@ export default function({ messages, template, types: t }) {
           right.isGenericType("Array") ||
           t.isArrayTypeAnnotation(right.getTypeAnnotation())
         ) {
-          return replaceWithArray(path);
+          replaceWithArray(path);
+          return;
         }
 
         let callback = spec;

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

@@ -24,20 +24,24 @@ export function insertBefore(nodes) {
     (this.parentPath.isForStatement() && this.key === "init")
   ) {
     if (this.node) nodes.push(this.node);
-    this.replaceExpressionWithStatements(nodes);
+    return this.replaceExpressionWithStatements(nodes);
   } else if (Array.isArray(this.container)) {
     return this._containerInsertBefore(nodes);
   } else if (this.isStatementOrBlock()) {
-    if (this.node) nodes.push(this.node);
+    const shouldInsertCurrentNode =
-    this.replaceWith(t.blockStatement(nodes));
+      this.node &&
+      (!this.isExpressionStatement() || this.node.expression != null);
+
+    this.replaceWith(
+      t.blockStatement(shouldInsertCurrentNode ? [this.node] : []),
+    );
+    return this.unshiftContainer("body", nodes);
   } else {
     throw new Error(
       "We don't know what to do with this node type. " +
         "We were previously a Statement but we can't fit in here?",
     );
   }
-
-  return [this];
 }
 
 export function _containerInsert(from, nodes) {
@@ -45,34 +49,14 @@ export function _containerInsert(from, nodes) {
 
   const paths = [];
 
+  this.container.splice(from, 0, ...nodes);
   for (let i = 0; i < nodes.length; i++) {
     const to = from + i;
-    const node = nodes[i];
+    const path = this.getSibling(`${to}`);
-    this.container.splice(to, 0, node);
-
-    if (this.context) {
-      const path = this.context.create(
-        this.parent,
-        this.container,
-        to,
-        this.listKey,
-      );
-
-      // While this path may have a context, there is currently no guarantee that the context
-      // will be the active context, because `popContext` may leave a final context in place.
-      // We should remove this `if` and always push once #4145 has been resolved.
-      if (this.context.queue) path.pushContext(this.context);
     paths.push(path);
-    } else {
+
-      paths.push(
+    if (this.context && this.context.queue) {
-        NodePath.get({
+      path.pushContext(this.context);
-          parentPath: this.parentPath,
-          parent: this.parent,
-          container: this.container,
-          listKey: this.listKey,
-          key: to,
-        }),
-      );
     }
   }
 
@@ -124,26 +108,24 @@ export function insertAfter(nodes) {
       );
       nodes.push(t.expressionStatement(temp));
     }
-    this.replaceExpressionWithStatements(nodes);
+    return this.replaceExpressionWithStatements(nodes);
   } else if (Array.isArray(this.container)) {
     return this._containerInsertAfter(nodes);
   } else if (this.isStatementOrBlock()) {
-    // Unshift current node if it's not an empty expression
+    const shouldInsertCurrentNode =
-    if (
       this.node &&
-      (!this.isExpressionStatement() || this.node.expression != null)
+      (!this.isExpressionStatement() || this.node.expression != null);
-    ) {
+
-      nodes.unshift(this.node);
+    this.replaceWith(
-    }
+      t.blockStatement(shouldInsertCurrentNode ? [this.node] : []),
-    this.replaceWith(t.blockStatement(nodes));
+    );
+    return this.pushContainer("body", nodes);
   } else {
     throw new Error(
       "We don't know what to do with this node type. " +
         "We were previously a Statement but we can't fit in here?",
     );
   }
-
-  return [this];
 }
 
 /**

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

@@ -50,13 +50,14 @@ export function replaceWithMultiple(nodes: Array<Object>) {
   t.inheritLeadingComments(nodes[0], this.node);
   t.inheritTrailingComments(nodes[nodes.length - 1], this.node);
   this.node = this.container[this.key] = null;
-  this.insertAfter(nodes);
+  const paths = this.insertAfter(nodes);
 
   if (this.node) {
     this.requeue();
   } else {
     this.remove();
   }
+  return paths;
 }
 
 /**
@@ -115,7 +116,7 @@ export function replaceWith(replacement) {
   }
 
   if (this.node === replacement) {
-    return;
+    return [this];
   }
 
   if (this.isProgram() && !t.isProgram(replacement)) {
@@ -136,6 +137,8 @@ export function replaceWith(replacement) {
     );
   }
 
+  let nodePath = "";
+
   if (this.isNodeType("Statement") && t.isExpression(replacement)) {
     if (
       !this.canHaveVariableDeclarationOrExpression() &&
@@ -144,6 +147,7 @@ export function replaceWith(replacement) {
     ) {
       // replacing a statement with an expression so wrap it in an expression statement
       replacement = t.expressionStatement(replacement);
+      nodePath = "expression";
     }
   }
 
@@ -172,6 +176,8 @@ export function replaceWith(replacement) {
 
   // requeue for visiting
   this.requeue();
+
+  return [nodePath ? this.get(nodePath) : this];
 }
 
 /**
@@ -206,8 +212,8 @@ export function replaceExpressionWithStatements(nodes: Array<Object>) {
   const toSequenceExpression = t.toSequenceExpression(nodes, this.scope);
 
   if (toSequenceExpression) {
-    this.replaceWith(toSequenceExpression);
+    return this.replaceWith(toSequenceExpression)[0].get("expressions");
-  } else {
+  }
   const container = t.arrowFunctionExpression([], t.blockStatement(nodes));
 
   this.replaceWith(t.callExpression(container, []));
@@ -241,10 +247,10 @@ export function replaceExpressionWithStatements(nodes: Array<Object>) {
     }
   }
 
-    this.get("callee").arrowFunctionToExpression();
+  const callee = this.get("callee");
+  callee.arrowFunctionToExpression();
 
-    return this.node;
+  return callee.get("body.body");
-  }
 }
 
 export function replaceInline(nodes: Object | Array<Object>) {
@@ -253,8 +259,9 @@ export function replaceInline(nodes: Object | Array<Object>) {
   if (Array.isArray(nodes)) {
     if (Array.isArray(this.container)) {
       nodes = this._verifyNodeList(nodes);
-      this._containerInsertAfter(nodes);
+      const paths = this._containerInsertAfter(nodes);
-      return this.remove();
+      this.remove();
+      return paths;
     } else {
       return this.replaceWithMultiple(nodes);
     }

packages/babel-traverse/test/modification.js

@@ -9,7 +9,7 @@ function getPath(code) {
   let path;
   traverse(ast, {
     Program: function(_path) {
-      path = _path;
+      path = _path.get("body.0");
       _path.stop();
     },
   });
@@ -18,22 +18,21 @@ function getPath(code) {
 }
 
 function generateCode(path) {
-  return generate(path.node).code;
+  return generate(path.parentPath.node).code;
 }
 
 describe("modification", function() {
   describe("pushContainer", function() {
     it("pushes identifier into params", function() {
       const rootPath = getPath("function test(a) {}");
-      const path = rootPath.get("body.0");
+      rootPath.pushContainer("params", t.identifier("b"));
-      path.pushContainer("params", t.identifier("b"));
 
       assert.equal(generateCode(rootPath), "function test(a, b) {}");
     });
 
     it("pushes identifier into block", function() {
       const rootPath = getPath("function test(a) {}");
-      const path = rootPath.get("body.0.body");
+      const path = rootPath.get("body");
       path.pushContainer("body", t.expressionStatement(t.identifier("b")));
 
       assert.equal(generateCode(rootPath), "function test(a) {\n  b;\n}");
@@ -42,18 +41,121 @@ describe("modification", function() {
   describe("unshiftContainer", function() {
     it("unshifts identifier into params", function() {
       const rootPath = getPath("function test(a) {}");
-      const path = rootPath.get("body.0");
+      rootPath.unshiftContainer("params", t.identifier("b"));
-      path.unshiftContainer("params", t.identifier("b"));
 
       assert.equal(generateCode(rootPath), "function test(b, a) {}");
     });
 
     it("unshifts identifier into block", function() {
       const rootPath = getPath("function test(a) {}");
-      const path = rootPath.get("body.0.body");
+      const path = rootPath.get("body");
       path.unshiftContainer("body", t.expressionStatement(t.identifier("b")));
 
       assert.equal(generateCode(rootPath), "function test(a) {\n  b;\n}");
     });
   });
+
+  describe("insertBefore", function() {
+    it("returns inserted path with BlockStatement", function() {
+      const rootPath = getPath("if (x) { y; }");
+      const path = rootPath.get("consequent.body.0");
+      const result = path.insertBefore(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[0].node, t.identifier("b"));
+      assert.equal(generateCode(rootPath), "if (x) {\n  b\n  y;\n}");
+    });
+
+    it("returns inserted path without BlockStatement", function() {
+      const rootPath = getPath("if (x) y;");
+      const path = rootPath.get("consequent");
+      const result = path.insertBefore(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[0].node, t.identifier("b"));
+      assert.equal(generateCode(rootPath), "if (x) {\n  b\n  y;\n}");
+    });
+
+    it("returns inserted path without BlockStatement without ExpressionStatement", function() {
+      const rootPath = getPath("if (x) for (var i = 0; i < 0; i++) {}");
+      const path = rootPath.get("consequent");
+      const result = path.insertBefore(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[result.length - 1].node, t.identifier("b"));
+      assert.equal(
+        generateCode(rootPath),
+        "if (x) {\n  b\n\n  for (var i = 0; i < 0; i++) {}\n}",
+      );
+    });
+
+    it("returns inserted path with BlockStatement without ExpressionStatement", function() {
+      const rootPath = getPath("if (x) { for (var i = 0; i < 0; i++) {} }");
+      const path = rootPath.get("consequent.body.0");
+      const result = path.insertBefore(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[result.length - 1].node, t.identifier("b"));
+      assert.equal(
+        generateCode(rootPath),
+        "if (x) {\n  b\n\n  for (var i = 0; i < 0; i++) {}\n}",
+      );
+    });
+  });
+
+  describe("insertAfter", function() {
+    it("returns inserted path with BlockStatement with ExpressionStatement", function() {
+      const rootPath = getPath("if (x) { y; }");
+      const path = rootPath.get("consequent.body.0");
+      const result = path.insertAfter(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[result.length - 1].node, t.identifier("b"));
+      assert.equal(generateCode(rootPath), "if (x) {\n  y;\n  b\n}");
+    });
+
+    it("returns inserted path without BlockStatement with ExpressionStatement", function() {
+      const rootPath = getPath("if (x) y;");
+      const path = rootPath.get("consequent");
+      const result = path.insertAfter(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[result.length - 1].node, t.identifier("b"));
+      assert.equal(generateCode(rootPath), "if (x) {\n  y;\n  b\n}");
+    });
+
+    it("returns inserted path without BlockStatement without ExpressionStatement", function() {
+      const rootPath = getPath("if (x) for (var i = 0; i < 0; i++) {}");
+      const path = rootPath.get("consequent");
+      const result = path.insertAfter(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[result.length - 1].node, t.identifier("b"));
+      assert.equal(
+        generateCode(rootPath),
+        "if (x) {\n  for (var i = 0; i < 0; i++) {}\n\n  b\n}",
+      );
+    });
+
+    it("returns inserted path with BlockStatement without ExpressionStatement", function() {
+      const rootPath = getPath("if (x) { for (var i = 0; i < 0; i++) {} }");
+      const path = rootPath.get("consequent.body.0");
+      const result = path.insertAfter(t.identifier("b"));
+
+      assert.equal(Array.isArray(result), true);
+      assert.equal(result.length, 1);
+      assert.deepEqual(result[result.length - 1].node, t.identifier("b"));
+      assert.equal(
+        generateCode(rootPath),
+        "if (x) {\n  for (var i = 0; i < 0; i++) {}\n\n  b\n}",
+      );
+    });
+  });
 });