improve `TreeWalker` performance (#3848)

1 files changed, 150 insertions, 129 deletions

diff --git a/lib/ast.js b/lib/ast.js
index 1f6f717d..4f56149c 100644
--- a/lib/ast.js
+++ b/lib/ast.js
@@ -110,11 +110,8 @@ var AST_Node = DEFNODE("Node", "start end", {
         start: "[AST_Token] The first token of this node",
         end: "[AST_Token] The last token of this node"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this);
-    },
     walk: function(visitor) {
-        return this._walk(visitor); // not sure the indirection will be any help
+        return visitor.visit(this);
     }
 }, null);
 
@@ -161,16 +158,17 @@ var AST_SimpleStatement = DEFNODE("SimpleStatement", "body", {
     $propdoc: {
         body: "[AST_Node] an expression node (should not be instanceof AST_Statement)"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.body._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.body.walk(visitor);
         });
     }
 }, AST_Statement);
 
 function walk_body(node, visitor) {
     node.body.forEach(function(node) {
-        node._walk(visitor);
+        node.walk(visitor);
     });
 }
 
@@ -179,9 +177,10 @@ var AST_Block = DEFNODE("Block", "body", {
     $propdoc: {
         body: "[AST_Statement*] an array of statements"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            walk_body(this, visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            walk_body(node, visitor);
         });
     }
 }, AST_Statement);
@@ -206,10 +205,11 @@ var AST_LabeledStatement = DEFNODE("LabeledStatement", "label", {
     $propdoc: {
         label: "[AST_Label] a label definition"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.label._walk(visitor);
-            this.body._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.label.walk(visitor);
+            node.body.walk(visitor);
         });
     },
     clone: function(deep) {
@@ -241,20 +241,22 @@ var AST_DWLoop = DEFNODE("DWLoop", "condition", {
 
 var AST_Do = DEFNODE("Do", null, {
     $documentation: "A `do` statement",
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.body._walk(visitor);
-            this.condition._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.body.walk(visitor);
+            node.condition.walk(visitor);
         });
     }
 }, AST_DWLoop);
 
 var AST_While = DEFNODE("While", null, {
     $documentation: "A `while` statement",
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.condition._walk(visitor);
-            this.body._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.condition.walk(visitor);
+            node.body.walk(visitor);
         });
     }
 }, AST_DWLoop);
@@ -266,12 +268,13 @@ var AST_For = DEFNODE("For", "init condition step", {
         condition: "[AST_Node?] the `for` termination clause, or null if empty",
         step: "[AST_Node?] the `for` update clause, or null if empty"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            if (this.init) this.init._walk(visitor);
-            if (this.condition) this.condition._walk(visitor);
-            if (this.step) this.step._walk(visitor);
-            this.body._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            if (node.init) node.init.walk(visitor);
+            if (node.condition) node.condition.walk(visitor);
+            if (node.step) node.step.walk(visitor);
+            node.body.walk(visitor);
         });
     }
 }, AST_IterationStatement);
@@ -282,11 +285,12 @@ var AST_ForIn = DEFNODE("ForIn", "init object", {
         init: "[AST_Node] the `for/in` initialization code",
         object: "[AST_Node] the object that we're looping through"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.init._walk(visitor);
-            this.object._walk(visitor);
-            this.body._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.init.walk(visitor);
+            node.object.walk(visitor);
+            node.body.walk(visitor);
         });
     }
 }, AST_IterationStatement);
@@ -296,10 +300,11 @@ var AST_With = DEFNODE("With", "expression", {
     $propdoc: {
         expression: "[AST_Node] the `with` expression"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
-            this.body._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
+            node.body.walk(visitor);
         });
     }
 }, AST_StatementWithBody);
@@ -380,13 +385,14 @@ var AST_Lambda = DEFNODE("Lambda", "name argnames uses_arguments length_read", {
         argnames: "[AST_SymbolFunarg*] array of function arguments",
         uses_arguments: "[boolean/S] tells whether this function accesses the arguments array"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            if (this.name) this.name._walk(visitor);
-            this.argnames.forEach(function(argname) {
-                argname._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            if (node.name) node.name.walk(visitor);
+            node.argnames.forEach(function(argname) {
+                argname.walk(visitor);
             });
-            walk_body(this, visitor);
+            walk_body(node, visitor);
         });
     }
 }, AST_Scope);
@@ -414,9 +420,10 @@ var AST_Exit = DEFNODE("Exit", "value", {
     $propdoc: {
         value: "[AST_Node?] the value returned or thrown by this statement; could be null for AST_Return"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, this.value && function() {
-            this.value._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            if (node.value) node.value.walk(visitor);
         });
     }
 }, AST_Jump);
@@ -434,9 +441,10 @@ var AST_LoopControl = DEFNODE("LoopControl", "label", {
     $propdoc: {
         label: "[AST_LabelRef?] the label, or null if none",
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, this.label && function() {
-            this.label._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            if (node.label) node.label.walk(visitor);
         });
     }
 }, AST_Jump);
@@ -457,11 +465,12 @@ var AST_If = DEFNODE("If", "condition alternative", {
         condition: "[AST_Node] the `if` condition",
         alternative: "[AST_Statement?] the `else` part, or null if not present"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.condition._walk(visitor);
-            this.body._walk(visitor);
-            if (this.alternative) this.alternative._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.condition.walk(visitor);
+            node.body.walk(visitor);
+            if (node.alternative) node.alternative.walk(visitor);
         });
     }
 }, AST_StatementWithBody);
@@ -473,10 +482,11 @@ var AST_Switch = DEFNODE("Switch", "expression", {
     $propdoc: {
         expression: "[AST_Node] the `switch` “discriminant”"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
-            walk_body(this, visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
+            walk_body(node, visitor);
         });
     }
 }, AST_Block);
@@ -494,10 +504,11 @@ var AST_Case = DEFNODE("Case", "expression", {
     $propdoc: {
         expression: "[AST_Node] the `case` expression"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
-            walk_body(this, visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
+            walk_body(node, visitor);
         });
     }
 }, AST_SwitchBranch);
@@ -510,11 +521,12 @@ var AST_Try = DEFNODE("Try", "bcatch bfinally", {
         bcatch: "[AST_Catch?] the catch block, or null if not present",
         bfinally: "[AST_Finally?] the finally block, or null if not present"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            walk_body(this, visitor);
-            if (this.bcatch) this.bcatch._walk(visitor);
-            if (this.bfinally) this.bfinally._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            walk_body(node, visitor);
+            if (node.bcatch) node.bcatch.walk(visitor);
+            if (node.bfinally) node.bfinally.walk(visitor);
         });
     }
 }, AST_Block);
@@ -524,10 +536,11 @@ var AST_Catch = DEFNODE("Catch", "argname", {
     $propdoc: {
         argname: "[AST_SymbolCatch] symbol for the exception"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.argname._walk(visitor);
-            walk_body(this, visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.argname.walk(visitor);
+            walk_body(node, visitor);
         });
     }
 }, AST_Block);
@@ -543,10 +556,11 @@ var AST_Definitions = DEFNODE("Definitions", "definitions", {
     $propdoc: {
         definitions: "[AST_VarDef*] array of variable definitions"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.definitions.forEach(function(defn) {
-                defn._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.definitions.forEach(function(defn) {
+                defn.walk(visitor);
             });
         });
     }
@@ -562,10 +576,11 @@ var AST_VarDef = DEFNODE("VarDef", "name value", {
         name: "[AST_SymbolVar] name of the variable",
         value: "[AST_Node?] initializer, or null of there's no initializer"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.name._walk(visitor);
-            if (this.value) this.value._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.name.walk(visitor);
+            if (node.value) node.value.walk(visitor);
         });
     }
 });
@@ -578,11 +593,12 @@ var AST_Call = DEFNODE("Call", "expression args", {
         expression: "[AST_Node] expression to invoke as function",
         args: "[AST_Node*] array of arguments"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
-            this.args.forEach(function(node) {
-                node._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
+            node.args.forEach(function(arg) {
+                arg.walk(visitor);
             });
         });
     }
@@ -597,10 +613,11 @@ var AST_Sequence = DEFNODE("Sequence", "expressions", {
     $propdoc: {
         expressions: "[AST_Node*] array of expressions (at least two)"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expressions.forEach(function(node) {
-                node._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expressions.forEach(function(expr) {
+                expr.walk(visitor);
             });
         });
     }
@@ -628,19 +645,21 @@ var AST_PropAccess = DEFNODE("PropAccess", "expression property", {
 
 var AST_Dot = DEFNODE("Dot", null, {
     $documentation: "A dotted property access expression",
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
         });
     }
 }, AST_PropAccess);
 
 var AST_Sub = DEFNODE("Sub", null, {
     $documentation: "Index-style property access, i.e. `a[\"foo\"]`",
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
-            this.property._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
+            node.property.walk(visitor);
         });
     }
 }, AST_PropAccess);
@@ -651,9 +670,10 @@ var AST_Unary = DEFNODE("Unary", "operator expression", {
         operator: "[string] the operator",
         expression: "[AST_Node] expression that this unary operator applies to"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.expression._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.expression.walk(visitor);
         });
     }
 });
@@ -673,10 +693,11 @@ var AST_Binary = DEFNODE("Binary", "operator left right", {
         operator: "[string] the operator",
         right: "[AST_Node] right-hand side expression"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.left._walk(visitor);
-            this.right._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.left.walk(visitor);
+            node.right.walk(visitor);
         });
     }
 });
@@ -688,11 +709,12 @@ var AST_Conditional = DEFNODE("Conditional", "condition consequent alternative",
         consequent: "[AST_Node]",
         alternative: "[AST_Node]"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.condition._walk(visitor);
-            this.consequent._walk(visitor);
-            this.alternative._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.condition.walk(visitor);
+            node.consequent.walk(visitor);
+            node.alternative.walk(visitor);
         });
     }
 });
@@ -708,10 +730,11 @@ var AST_Array = DEFNODE("Array", "elements", {
     $propdoc: {
         elements: "[AST_Node*] array of elements"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.elements.forEach(function(element) {
-                element._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.elements.forEach(function(element) {
+                element.walk(visitor);
             });
         });
     }
@@ -722,10 +745,11 @@ var AST_Object = DEFNODE("Object", "properties", {
     $propdoc: {
         properties: "[AST_ObjectProperty*] array of properties"
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.properties.forEach(function(prop) {
-                prop._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.properties.forEach(function(prop) {
+                prop.walk(visitor);
             });
         });
     }
@@ -737,9 +761,10 @@ var AST_ObjectProperty = DEFNODE("ObjectProperty", "key value", {
         key: "[string|AST_SymbolAccessor] property name. For ObjectKeyVal this is a string. For getters and setters this is an AST_SymbolAccessor.",
         value: "[AST_Node] property value.  For getters and setters this is an AST_Accessor."
     },
-    _walk: function(visitor) {
-        return visitor._visit(this, function() {
-            this.value._walk(visitor);
+    walk: function(visitor) {
+        var node = this;
+        return visitor.visit(node, function() {
+            node.value.walk(visitor);
         });
     }
 });
@@ -891,21 +916,17 @@ var AST_True = DEFNODE("True", null, {
 /* -----[ TreeWalker ]----- */
 
 function TreeWalker(callback) {
-    this.visit = callback;
-    this.stack = [];
+    this.callback = callback;
     this.directives = Object.create(null);
+    this.stack = [];
 }
 TreeWalker.prototype = {
-    _visit: function(node, descend) {
+    visit: function(node, descend) {
         this.push(node);
-        var ret = this.visit(node, descend ? function() {
-            descend.call(node);
-        } : noop);
-        if (!ret && descend) {
-            descend.call(node);
-        }
+        var done = this.callback(node, descend || noop);
+        if (!done && descend) descend();
         this.pop();
-        return ret;
+        return done;
     },
     parent: function(n) {
         return this.stack[this.stack.length - 2 - (n || 0)];