Refactor traversal to avoid closures

lib/6to5/traverse/index.js

@@ -4,17 +4,101 @@ var Scope = require("./scope");
 var t     = require("../types");
 var _     = require("lodash");
 
+function TraversalContext(previousContext) {
+  this.didSkip = false;
+  this.didRemove = false;
+  this.didStop = false;
+  this.didFlatten = previousContext ? previousContext.didFlatten : false;
+}
+
+TraversalContext.prototype.flatten = function () {
+  this.didFlatten = true;
+};
+
+TraversalContext.prototype.remove = function () {
+  this.didRemove = true;
+  this.skip();
+};
+
+TraversalContext.prototype.skip = function () {
+  this.didSkip = true;
+};
+
+TraversalContext.prototype.stop = function () {
+  this.didStop = true;
+  this.skip();
+};
+
+TraversalContext.prototype.maybeReplace = function (result, obj, key, node) {
+  if (result === false) return node;
+  if (result == null) return node;
+
+  var isArray = _.isArray(result);
+
+  // inherit comments from original node to the first replacement node
+  var inheritTo = result;
+  if (isArray) inheritTo = result[0];
+  if (inheritTo) t.inheritsComments(inheritTo, node);
+
+  // replace the node
+  node = obj[key] = result;
+
+  // 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();
+  }
+
+  return node;
+};
+
+TraversalContext.prototype.visit = function (obj, key, opts, scope, parent) {
+  var node = obj[key];
+  if (!node) return;
+
+  // type is blacklisted
+  if (opts.blacklist && opts.blacklist.indexOf(node.type) > -1) return;
+
+  var result;
+  var ourScope = scope;
+  if (t.isScope(node)) ourScope = new Scope(node, scope);
+
+  // enter
+  if (opts.enter) {
+    result = opts.enter.call(this, node, parent, ourScope);
+    node = this.maybeReplace(result, obj, key, node);
+
+    if (this.didRemove) {
+      obj[key] = null;
+      this.flatten();
+    }
+
+    // stop traversal
+    if (this.didSkip) return;
+  }
+
+  // traverse node
+  traverse(node, opts, ourScope);
+
+  // exit
+  if (opts.exit) {
+    result = opts.exit.call(this, node, parent, ourScope);
+    node = this.maybeReplace(result, obj, key, node);
+  }
+};
+
 function traverse(parent, opts, scope) {
   // falsy node
   if (!parent) return;
 
-  var i, j;
-
   // array of nodes
   if (_.isArray(parent)) {
-    for (i = 0; i < parent.length; i++) {
+    for (var i = 0; i < parent.length; i++)
       traverse(parent[i], opts, scope);
-    }
     return;
   }
 
@@ -23,99 +107,21 @@ function traverse(parent, opts, scope) {
   if (!keys) return;
 
   opts = opts || {};
+  var context = null;
 
-  var stopped = false;
-
-  for (i = 0; i < keys.length; i++) {
-    var key = keys[i];
+  for (var j = 0; j < keys.length; j++) {
+    var key = keys[j];
     var nodes = parent[key];
     if (!nodes) continue;
 
-    var flatten = false;
-
-    var handle = function (obj, key) {
-      var node = obj[key];
-      if (!node) return;
-
-      // type is blacklisted
-      if (opts.blacklist && opts.blacklist.indexOf(node.type) > -1) return;
-
-      // replace node
-      var maybeReplace = function (result) {
-        if (result === false) return;
-        if (result == null) return;
-
-        var isArray = _.isArray(result);
-
-        // inherit comments from original node to the first replacement node
-        var inheritTo = result;
-        if (isArray) inheritTo = result[0];
-        if (inheritTo) t.inheritsComments(inheritTo, node);
-
-        // replace the node
-        node = obj[key] = result;
-
-        if (isArray) flatten = true;
-
-        // 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);
-        }
-      };
-
-      var skipped = false;
-      var removed = false;
-
-      var context = {
-        stop: function () {
-          skipped = stopped = true;
-        },
-
-        skip: function () {
-          skipped = true;
-        },
-
-        remove: function () {
-          this.skip();
-          removed = true;
-        }
-      };
-
-      //
-      var ourScope = scope;
-      if (t.isScope(node)) ourScope = new Scope(node, scope);
-
-      // enter
-      if (opts.enter) {
-        var result = opts.enter.call(context, node, parent, ourScope);
-        maybeReplace(result);
-
-        if (removed) {
-          obj[key] = null;
-          flatten = true;
-        }
-
-        // stop iteration
-        if (skipped) return;
-      }
-
-      // traverse node
-      traverse(node, opts, ourScope);
-
-      // exit
-      if (opts.exit) {
-        maybeReplace(opts.exit.call(context, node, parent, ourScope));
-      }
-    };
-
     if (_.isArray(nodes)) {
-      for (j = 0; j < nodes.length; j++) {
-        handle(nodes, j);
-        if (stopped) return;
+      for (var k = 0; k < nodes.length; k++) {
+        context = new TraversalContext(context);
+        context.visit(nodes, k, opts, scope, parent);
+        if (context.didStop) return;
       }
 
-      if (flatten) {
+      if (context && context.didFlatten) {
         parent[key] = _.flatten(parent[key]);
 
         if (key === "body") {
@@ -124,8 +130,9 @@ function traverse(parent, opts, scope) {
         }
       }
     } else {
-      handle(parent, key);
-      if (stopped) return;
+      context = new TraversalContext(context);
+      context.visit(parent, key, opts, scope, parent);
+      if (context.didStop) return;
     }
   }
 }