split up traversal so it's easier to maintain and extend later on

lib/6to5/traversal/context.js

@@ -0,0 +1,78 @@
+"use strict";
+
+/* jshint maxparams:7 */
+
+module.exports = TraversalContext;
+
+var TraversalIteration = require("./iteration");
+var flatten            = require("lodash/array/flatten");
+var compact            = require("lodash/array/compact");
+
+function TraversalContext(scope, opts, state) {
+  this.shouldFlatten = false;
+
+  this.scope = scope;
+  this.state = state;
+  this.opts  = opts;
+
+  this.reset();
+}
+
+TraversalContext.prototype.flatten = function () {
+  this.shouldFlatten = true;
+};
+
+TraversalContext.prototype.remove = function () {
+  this.shouldRemove = true;
+  this.shouldSkip   = true;
+};
+
+TraversalContext.prototype.skip = function () {
+  this.shouldSkip = true;
+};
+
+TraversalContext.prototype.stop = function () {
+  this.shouldStop = true;
+  this.shouldSkip = true;
+};
+
+TraversalContext.prototype.reset = function () {
+  this.shouldRemove = false;
+  this.shouldSkip   = false;
+  this.shouldStop   = false;
+};
+
+TraversalContext.prototype.visitNode = function (node, obj, key) {
+  this.reset();
+  var iteration = new TraversalIteration(this, node, obj, key);
+  return iteration.visit();
+};
+
+TraversalContext.prototype.visit = function (node, key) {
+  var nodes = node[key];
+  if (!nodes) return;
+
+  if (!Array.isArray(nodes)) {
+    return this.visitNode(node, node, key);
+  }
+
+  // nothing to traverse!
+  if (nodes.length === 0) {
+    return;
+  }
+
+  for (var i = 0; i < nodes.length; i++) {
+    if (nodes[i] && this.visitNode(node, nodes, i)) {
+      return true;
+    }
+  }
+
+  if (this.shouldFlatten) {
+    node[key] = flatten(node[key]);
+
+    if (key === "body") {
+      // we can safely compact this
+      node[key] = compact(node[key]);
+    }
+  }
+};

lib/6to5/traversal/index.js

@@ -2,174 +2,9 @@
 
 module.exports = traverse;
 
-/* jshint maxparams:7 */
+var TraversalContext = require("./context");
-
-var Scope    = require("./scope");
-var t        = require("../types");
 var contains         = require("lodash/collection/contains");
-var flatten  = require("lodash/array/flatten");
+var t                = require("../types");
-var compact  = require("lodash/array/compact");
-
-function TraversalContext(scope) {
-  this.shouldFlatten = false;
-  this.shouldRemove  = false;
-  this.shouldSkip    = false;
-  this.shouldStop    = false;
-  this.scope         = scope;
-}
-
-TraversalContext.prototype.flatten = function () {
-  this.shouldFlatten = true;
-};
-
-TraversalContext.prototype.remove = function () {
-  this.shouldRemove = true;
-  this.shouldSkip   = true;
-};
-
-TraversalContext.prototype.skip = function () {
-  this.shouldSkip = true;
-};
-
-TraversalContext.prototype.stop = function () {
-  this.shouldStop = true;
-  this.shouldSkip = true;
-};
-
-TraversalContext.prototype.reset = function () {
-  this.shouldRemove = false;
-  this.shouldSkip   = false;
-  this.shouldStop   = false;
-};
-
-TraversalContext.prototype.maybeRemove = function (obj, key) {
-  if (this.shouldRemove) {
-    obj[key] = null;
-    this.flatten();
-  }
-};
-
-TraversalContext.prototype.replaceNode = function (obj, key, node, replacement, scope) {
-  var isArray = Array.isArray(replacement);
-
-  // inherit comments from original node to the first replacement node
-  var inheritTo = replacement;
-  if (isArray) inheritTo = replacement[0];
-  if (inheritTo) t.inheritsComments(inheritTo, node);
-
-  // replace the node
-  obj[key] = replacement;
-
-  var file = this.scope && this.scope.file;
-  if (file) {
-    if (isArray) {
-      for (var i = 0; i < replacement.length; i++) {
-        file.checkNode(replacement[i], scope);
-      }
-    } else {
-      file.checkNode(replacement, scope);
-    }
-  }
-
-  // we're replacing a statement or block node with an array of statements so we better
-  // ensure that it's a block
-  if (isArray && contains(t.STATEMENT_OR_BLOCK_KEYS, key) && !t.isBlockStatement(obj)) {
-    t.ensureBlock(obj, key);
-  }
-
-  if (isArray) {
-    this.flatten();
-  }
-};
-
-TraversalContext.prototype.call = function (fn, obj, key, node, parent, scope, state) {
-  var replacement = fn(node, parent, scope, this, state);
-
-  if (replacement) {
-    this.replaceNode(obj, key, node, replacement, scope);
-    node = replacement;
-  }
-
-  this.maybeRemove(obj, key);
-
-  return node;
-};
-
-TraversalContext.prototype.visitNode = function (obj, key, opts, scope, parent, state) {
-  this.reset();
-
-  var node = obj[key];
-
-  // type is blacklisted
-  if (opts.blacklist && opts.blacklist.indexOf(node.type) > -1) {
-    return;
-  }
-
-  var ourScope = scope;
-  // we're entering a new scope so let's construct it!
-  if (!opts.noScope && t.isScope(node)) {
-    ourScope = new Scope(node, parent, scope);
-  }
-
-  node = this.call(opts.enter, obj, key, node, parent, ourScope, state);
-
-  if (this.shouldSkip) {
-    return this.shouldStop;
-  }
-
-  if (Array.isArray(node)) {
-    // traverse over these replacement nodes we purposely don't call exitNode
-    // as the original node has been destroyed
-    for (var i = 0; i < node.length; i++) {
-      traverseNode(node[i], opts, ourScope, state);
-    }
-  } else {
-    traverseNode(node, opts, ourScope, state);
-    this.call(opts.exit, obj, key, node, parent, ourScope, state);
-  }
-
-  return this.shouldStop;
-};
-
-TraversalContext.prototype.visit = function (node, key, opts, scope, state) {
-  var nodes = node[key];
-  if (!nodes) return;
-
-  if (!Array.isArray(nodes)) {
-    return this.visitNode(node, key, opts, scope, node, state);
-  }
-
-  if (nodes.length === 0) {
-    return;
-  }
-
-  for (var i = 0; i < nodes.length; i++) {
-    if (nodes[i] && this.visitNode(nodes, i, opts, scope, node, state)) {
-      return true;
-    }
-  }
-
-  if (this.shouldFlatten) {
-    node[key] = flatten(node[key]);
-
-    if (key === "body") {
-      // we can safely compact this
-      node[key] = compact(node[key]);
-    }
-  }
-};
-
-function traverseNode(node, opts, scope, state) {
-  var keys = t.VISITOR_KEYS[node.type];
-  if (!keys) return;
-
-  var context = new TraversalContext(scope);
-  for (var i = 0; i < keys.length; i++) {
-    if (context.visit(node, keys[i], opts, scope, state)) {
-      return;
-    }
-  }
-}
 
 function traverse(parent, opts, scope, state) {
   if (!parent) return;
@@ -187,13 +22,25 @@ function traverse(parent, opts, scope, state) {
   // array of nodes
   if (Array.isArray(parent)) {
     for (var i = 0; i < parent.length; i++) {
-      traverseNode(parent[i], opts, scope, state);
+      traverse.node(parent[i], opts, scope, state);
     }
   } else {
-    traverseNode(parent, opts, scope, state);
+    traverse.node(parent, opts, scope, state);
   }
 }
 
+traverse.node = function (node, opts, scope, state) {
+  var keys = t.VISITOR_KEYS[node.type];
+  if (!keys) return;
+
+  var context = new TraversalContext(scope, opts, state);
+  for (var i = 0; i < keys.length; i++) {
+    if (context.visit(node, keys[i])) {
+      return;
+    }
+  }
+};
+
 function clearNode(node) {
   node._declarations = null;
   node.extendedRange = null;

lib/6to5/traversal/iteration.js

@@ -0,0 +1,126 @@
+"use strict";
+
+module.exports = TraversalIteration;
+
+/* jshint maxparams:7 */
+
+var traverse = require("./index");
+var contains = require("lodash/collection/contains");
+var Scope    = require("./scope");
+var t        = require("../types");
+
+function TraversalIteration(context, parent, obj, key) {
+  this.key     = key;
+  this.obj     = obj;
+  this.context = context;
+  this.parent  = parent;
+  this.scope   = TraversalIteration.getScope(this.getNode(), parent, context.scope);
+  this.state   = context.state;
+}
+
+TraversalIteration.getScope = function (node, parent, scope) {
+  var ourScope = scope;
+
+  // we're entering a new scope so let's construct it!
+  if (t.isScope(node)) {
+    ourScope = new Scope(node, parent, scope);
+  }
+
+  return ourScope;
+};
+
+TraversalIteration.prototype.maybeRemove = function () {
+  if (this.context.shouldRemove) {
+    this.setNode(null);
+    this.context.flatten();
+  }
+};
+
+TraversalIteration.prototype.setNode = function (val) {
+  return this.obj[this.key] = val;
+};
+
+TraversalIteration.prototype.getNode = function () {
+  return this.obj[this.key];
+};
+
+TraversalIteration.prototype.replaceNode = function (replacement, scope) {
+  var isArray = Array.isArray(replacement);
+
+  // inherit comments from original node to the first replacement node
+  var inheritTo = replacement;
+  if (isArray) inheritTo = replacement[0];
+  if (inheritTo) t.inheritsComments(inheritTo, this.getNode());
+
+  // replace the node
+  this.setNode(replacement);
+
+  var file = this.scope && this.scope.file;
+  if (file) {
+    if (isArray) {
+      for (var i = 0; i < replacement.length; i++) {
+        file.checkNode(replacement[i], scope);
+      }
+    } else {
+      file.checkNode(replacement, scope);
+    }
+  }
+
+  // we're replacing a statement or block node with an array of statements so we better
+  // ensure that it's a block
+  if (isArray) {
+    if (contains(t.STATEMENT_OR_BLOCK_KEYS, this.key) && !t.isBlockStatement(this.obj)) {
+      t.ensureBlock(this.obj, this.key);
+    }
+
+    this.context.flatten();
+  }
+};
+
+TraversalIteration.prototype.call = function (fn) {
+  var node = this.getNode();
+  var replacement = fn.call(this, node, this.parent, this.scope, this.context, this.state);
+
+  if (replacement) {
+    this.replaceNode(replacement);
+    node = replacement;
+  }
+
+  this.maybeRemove();
+
+  return node;
+};
+
+TraversalIteration.prototype.visit = function () {
+  this.context.reset();
+
+  var state = this.context.state;
+  var opts  = this.context.opts;
+  var node  = this.getNode();
+
+  // type is blacklisted
+  if (opts.blacklist && opts.blacklist.indexOf(node.type) > -1) {
+    return;
+  }
+
+  this.call(opts.enter);
+
+  if (this.context.shouldSkip) {
+    return this.context.shouldStop;
+  }
+
+  node = this.getNode();
+
+  if (Array.isArray(node)) {
+    // traverse over these replacement nodes we purposely don't call exitNode
+    // as the original node has been destroyed
+    for (var i = 0; i < node.length; i++) {
+      traverse.node(node[i], opts, this.scope, state);
+    }
+  } else {
+    traverse.node(node, opts, this.scope, state);
+    this.call(opts.exit);
+  }
+
+  return this.context.shouldStop;
+};