impose a strict API for traversal paths and ensure AST doesn't have holes mid-transform

src/babel/transformation/transformers/es6/block-scoping-tdz.js

@@ -3,6 +3,7 @@ import * as t from "../../../types";
 var visitor = {
   enter(node, parent, scope, state) {
     if (!this.isReferencedIdentifier()) return;
+    if (t.isFor(parent) && parent.left === node) return;
 
     var declared = state.letRefs[node.name];
     if (!declared) return;
@@ -19,7 +20,7 @@ var visitor = {
 
     if (t.isAssignmentExpression(parent) || t.isUpdateExpression(parent)) {
       if (parent._ignoreBlockScopingTDZ) return;
-      this.parentPath.node = t.sequenceExpression([assert, parent]);
+      this.parentPath.replaceWith(t.sequenceExpression([assert, parent]));
     } else {
       return t.logicalExpression("&&", assert, node);
     }

src/babel/transformation/transformers/es6/for-of.js

@@ -33,7 +33,7 @@ export function ForOfStatement(node, parent, scope, file) {
 
   t.inherits(loop, node);
 
-  if (build.replaceParent) this.parentPath.node = build.node;
+  if (build.replaceParent) this.parentPath.replaceWith(build.node);
   return build.node;
 }
 

src/babel/transformation/transformers/es6/parameters.rest.js

@@ -94,7 +94,7 @@ exports.Function = function (node, parent, scope, file) {
   if (state.canOptimise && state.candidates.length) {
     for (var i = 0; i < state.candidates.length; i++) {
       var candidate = state.candidates[i];
-      candidate.node = argsId;
+      candidate.replaceWith(argsId);
       optimizeMemberExpression(candidate.parent, node.params.length);
     }
     return;

src/babel/transformation/transformers/internal/block-hoist.js

@@ -19,7 +19,7 @@ export var BlockStatement = {
     if (!hasChange) return;
 
     var nodePriorities = groupBy(node.body, function (bodyNode) {
-      var priority = bodyNode._blockHoist;
+      var priority = bodyNode && bodyNode._blockHoist;
       if (priority == null) priority = 1;
       if (priority === true) priority = 2;
       return priority;

src/babel/traversal/context.js

@@ -1,24 +1,17 @@
 import TraversalPath from "./path";
-import flatten from "lodash/array/flatten";
 import compact from "lodash/array/compact";
 import * as t from "../types";
 
 export default class TraversalContext {
   constructor(scope, opts, state, parentPath) {
-    this.shouldFlatten = false;
+    this.parentPath = parentPath;
-    this.parentPath    = parentPath;
+    this.scope      = scope;
-    this.scope         = scope;
+    this.state      = state;
-    this.state         = state;
+    this.opts       = opts;
-    this.opts          = opts;
   }
 
-  flatten() {
+  create(node, obj, key) {
-    this.shouldFlatten = true;
+    return TraversalPath.get(this.parentPath, this, node, obj, key);
-  }
-
-  visitNode(node, obj, key) {
-    var iteration = TraversalPath.get(this.parentPath, this, node, obj, key);
-    return iteration.visit();
   }
 
   visit(node, key) {
@@ -26,7 +19,7 @@ export default class TraversalContext {
     if (!nodes) return;
 
     if (!Array.isArray(nodes)) {
-      return this.visitNode(node, node, key);
+      return this.create(node, node, key).visit();
     }
 
     // nothing to traverse!
@@ -34,19 +27,15 @@ export default class TraversalContext {
       return;
     }
 
-    // todo: handle nodes popping in and out of existence
+    var queue = [];
-    for (var i = 0; i < nodes.length; i++) {
+
-      if (nodes[i] && this.visitNode(node, nodes, i)) {
+    for (let i = 0; i < nodes.length; i++) {
-        return true;
+      if (nodes[i]) queue.push(this.create(node, nodes, i));
-      }
     }
 
-    if (this.shouldFlatten) {
+    for (let i = 0; i < queue.length; i++) {
-      node[key] = flatten(node[key]);
+      if (queue[i].visit()) {
-
+        return true;
-      if (t.FLATTENABLE_KEYS.indexOf(key) >= 0) {
-        // we can safely compact this
-        node[key] = compact(node[key]);
       }
     }
   }

src/babel/traversal/path/index.js

@@ -78,46 +78,50 @@ export default class TraversalPath {
     return ourScope;
   }
 
-  getParentArrayPath() {
-    var path = this;
-    while (!Array.isArray(path.container)) {
-      path = path.parentPath;
-    }
-    return path;
-  }
-
   insertBefore(nodes) {
-    this.getParentArrayPath()._insertBefore(nodes);
+    this.checkNodes(nodes);
+
+    if (this.isPreviousType("Expression")) {
+      if (this.node) nodes.push(this.node);
+      this.replaceExpressionWithStatements(nodes);
+    } else {
+      throw new Error("no idea what to do with this");
+    }
   }
 
   insertAfter(nodes) {
-    this.getParentArrayPath()._insertAfter(nodes);
+    this.checkNodes(nodes);
-  }
 
-  _insertBefore(nodes) {
+    if (this.isPreviousType("Statement")) {
-    throw new Error("to be implemented");
+      if (Array.isArray(this.container)) {
-  }
+        for (var i = 0; i < nodes.length; i++) {
-
+          this.container.splice(this.key + 1 + i, 0, nodes[i]);
-  _insertAfter(nodes) {
+        }
-    if (this.isStatement()) {
+        this.updateSiblingKeys(this.key + nodes.length, nodes.length);
-      for (var i = 0; i < nodes.length; i++) {
+      } else if (includes(t.STATEMENT_OR_BLOCK_KEYS, this.key) && !t.isBlockStatement(this.container)) {
-        let key = this.key + 1 + i;
+        this.container[this.key] = t.blockStatement(nodes);
-        this.container.splice(key, 0, nodes[i]);
+      } else {
+        throw new Error("no idea what to do with this");
       }
-      this.incrementSiblingPaths(this.key + nodes.length, nodes.length);
+    } else if (this.isPreviousType("Expression")) {
+      if (this.node) {
+        var temp = this.scope.generateTemp();
+        nodes.unshift(t.expressionStatement(t.assignmentExpression("=", temp, this.node)));
+        nodes.push(t.expressionStatement(temp));
+      }
+      this.replaceExpressionWithStatements(nodes);
     } else {
-      let key = this.key + 1;
+      throw new Error("no idea what to do with this");
-      this.container.splice(key, 0, null);
-      this.incrementSiblingPaths(key, 1);
-      this.getSibling(key).setStatementsToExpression(nodes);
     }
   }
 
-  incrementSiblingPaths(fromIndex, incrementBy) {
+  updateSiblingKeys(fromIndex, incrementBy) {
     var paths = this.container._paths;
-    for (var i = 0; i > paths.length; i++) {
+    for (var i = 0; i < paths.length; i++) {
-      let path = paths[path];
+      let path = paths[i];
-      if (path.key >= fromIndex) path.key += incrementBy;
+      if (path.key >= fromIndex) {
+        path.key += incrementBy;
+      }
     }
   }
 
@@ -138,6 +142,7 @@ export default class TraversalPath {
   setContext(parentPath, context, key, file?) {
     this.shouldSkip = false;
     this.shouldStop = false;
+    this.removed    = false;
 
     this.parentPath = parentPath || this.parentPath;
     this.key        = key;
@@ -153,10 +158,18 @@ export default class TraversalPath {
     this.type = this.node && this.node.type;
   }
 
+  _remove() {
+    if (Array.isArray(this.container)) {
+      this.container.splice(this.key, 1);
+      this.updateSiblingKeys(this.key, -1);
+    } else {
+      this.container[this.key] = null;
+    }
+  }
+
   remove() {
-    this._refresh(this.node, []);
+    this._remove();
-    this.container[this.key] = null;
+    this.removed = true;
-    this.flatten();
   }
 
   skip() {
@@ -168,19 +181,6 @@ export default class TraversalPath {
     this.shouldSkip = true;
   }
 
-  flatten() {
-    this.context.flatten();
-  }
-
-  _refresh(oldNode, newNodes) {
-    // todo
-  }
-
-  refresh() {
-    var node = this.node;
-    this._refresh(node, [node]);
-  }
-
   errorWithNode(msg, Error = SyntaxError) {
     var loc = this.node.loc.start;
     var err = new Error(`Line ${loc.line}: ${msg}`);
@@ -189,31 +189,58 @@ export default class TraversalPath {
   }
 
   get node() {
-    return this.container[this.key];
+    if (this.removed) {
+      return null;
+    } else {
+      return this.container[this.key];
+    }
   }
 
   set node(replacement) {
-    if (!replacement) return this.remove();
+    throw new Error("Don't use `path.node = newNode;`, use `path.replaceWith(newNode)` or `path.replaceWithMultiple([newNode])`");
+  }
 
-    var oldNode = this.node;
+  replaceWithMultiple(nodes: Array<Object>) {
-
+    if (nodes.indexOf(this.node) >= 0) {
-    var isArray = Array.isArray(replacement);
+      // todo: check for inclusion of current node in `nodes` and yell at the user if it's in there and tell them to use `insertBefore` or `insertAfter`
-    if (isArray && replacement.length === 1) {
-      isArray = false;
-      replacement = replacement[0];
     }
 
-    var replacements = isArray ? replacement : [replacement];
+    for (var i = 0; i < nodes.length; i++) {
-
+      var node = nodes[i];
-    // inherit comments from original node to the first replacement node
+      if (!node) throw new Error(`Falsy node passed to \`path.replaceWithMultiple()\` with the index of ${i}`);
-    var inheritTo = replacements[0];
+    }
-    if (inheritTo && oldNode) t.inheritsComments(inheritTo, oldNode);
 
     //
-    if (t.isStatement(replacements[0]) && t.isType(this.type, "Expression")) {
+
-      return this.setStatementsToExpression(replacements);
+    this.container[this.key] = null;
+    this.insertAfter(nodes);
+    if (!this.node) this.remove();
+  }
+
+  replaceWith(replacement, arraysAllowed) {
+    if (this.removed) {
+      throw new Error("Trying to replace a node that we've removed");
     }
 
+    if (!replacement) {
+      throw new Error("You passed `path.replaceWith()` a falsy node, use `path.remove()` instead");
+    }
+
+    if (Array.isArray(replacement)) {
+      if (arraysAllowed) {
+        return this.replaceWithMultiple(replacement);
+      } else {
+        throw new Error("Don't use `path.replaceWith()` with an array of nodes, use `path.replaceWithMultiple()`");
+      }
+    }
+
+    if (this.isPreviousType("Expression") && t.isStatement(replacement)) {
+      return this.replaceExpressionWithStatements([replacement]);
+    }
+
+    var oldNode = this.node;
+    if (oldNode) t.inheritsComments(replacement, oldNode);
+
     // replace the node
     this.container[this.key] = replacement;
     this.type = replacement.type;
@@ -221,21 +248,16 @@ export default class TraversalPath {
     // potentially create new scope
     this.setScope();
 
-    var file = this.scope && this.scope.file;
+    this.checkNodes([replacement]);
-    if (file) {
+  }
-      for (var i = 0; i < replacements.length; i++) {
-        file.checkNode(replacements[i], this.scope);
-      }
-    }
 
-    // we're replacing a statement or block node with an array of statements so we better
+  checkNodes(nodes) {
-    // ensure that it's a block
+    var scope = this.scope;
-    if (isArray) {
+    var file  = scope && scope.file;
-      if (includes(t.STATEMENT_OR_BLOCK_KEYS, this.key) && !t.isBlockStatement(this.container)) {
+    if (!file) return;
-        t.ensureBlock(this.container, this.key);
-      }
 
-      this.flatten();
+    for (var i = 0; i < nodes.length; i++) {
+      file.checkNode(nodes[i], scope);
     }
   }
 
@@ -260,11 +282,11 @@ export default class TraversalPath {
     return paths;
   }
 
-  setStatementsToExpression(nodes: Array) {
+  replaceExpressionWithStatements(nodes: Array) {
     var toSequenceExpression = t.toSequenceExpression(nodes, this.scope);
 
     if (toSequenceExpression) {
-      return this.node = toSequenceExpression;
+      return this.replaceWith(toSequenceExpression);
     } else {
       var container = t.functionExpression(null, [], t.blockStatement(nodes));
       container.shadow = true;
@@ -274,11 +296,11 @@ export default class TraversalPath {
       for (var i = 0; i < last.length; i++) {
         var lastNode = last[i];
         if (lastNode.isExpressionStatement()) {
-          lastNode.node = t.returnStatement(lastNode.node.expression);
+          lastNode.replaceWith(t.returnStatement(lastNode.node.expression));
         }
       }
 
-      this.node = t.callExpression(container, []);
+      this.replaceWith(t.callExpression(container, []));
 
       this.traverse(hoistVariablesVisitor);
 
@@ -295,10 +317,7 @@ export default class TraversalPath {
     if (opts[node.type]) fn = opts[node.type][key] || fn;
 
     var replacement = fn.call(this, node, this.parent, this.scope, this.state);
-
+    if (replacement) this.replaceWith(replacement, true);
-    if (replacement) {
-      this.node = replacement;
-    }
   }
 
   isBlacklisted(): boolean {
@@ -513,6 +532,10 @@ export default class TraversalPath {
     return t.isScope(this.node, this.parent);
   }
 
+  isPreviousType(type: string): boolean {
+    return t.isType(this.type, type);
+  }
+
   isTypeGeneric(genericName: string, opts = {}): boolean {
     var typeInfo = this.getTypeAnnotation();
     var type     = typeInfo.annotation;