From 5e7f5f390ad6f2fec3270fa829b31a7b229999bb Mon Sep 17 00:00:00 2001
From: Jeremy Ashkenas <jashkenas@gmail.com>
Date: Sat, 27 Feb 2010 01:22:21 -0500
Subject: [PATCH] adding a traverse method to the AST, so we can do fancy
 processing from external scripts.

---
 lib/nodes.js     | 35 +++++++++++++++++++++++++++++++----
 src/nodes.coffee | 19 ++++++++++++++++---
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/lib/nodes.js b/lib/nodes.js
index 51b22bfe..0b1f7125 100644
--- a/lib/nodes.js
+++ b/lib/nodes.js
@@ -142,12 +142,27 @@
       if (block(node)) {
         return true;
       }
-      if (node instanceof BaseNode && node.contains(block)) {
+      if (node.contains && node.contains(block)) {
         return true;
       }
     }
     return false;
   };
+  // Perform an in-order traversal of the AST.
+  BaseNode.prototype.traverse = function traverse(block) {
+    var _a, _b, _c, _d, node;
+    _a = []; _b = this.children;
+    for (_c = 0, _d = _b.length; _c < _d; _c++) {
+      node = _b[_c];
+      _a.push((function() {
+        block(node);
+        if (node.traverse) {
+          return node.traverse(block);
+        }
+      }).call(this));
+    }
+    return _a;
+  };
   // toString representation of the node, for inspecting the parse tree.
   BaseNode.prototype.toString = function toString(idt) {
     var _a, _b, _c, _d, child;
@@ -854,12 +869,24 @@
     top_sensitive: function top_sensitive() {
       return true;
     },
+    real_children: function real_children() {
+      return flatten([this.params, this.body.expressions]);
+    },
+    traverse: function traverse(block) {
+      var _a, _b, _c, _d, child;
+      block(this);
+      _a = []; _b = this.real_children();
+      for (_c = 0, _d = _b.length; _c < _d; _c++) {
+        child = _b[_c];
+        _a.push(block(child));
+      }
+      return _a;
+    },
     toString: function toString(idt) {
-      var _a, _b, _c, _d, child, children;
+      var _a, _b, _c, _d, child;
       idt = idt || '';
-      children = flatten([this.params, this.body.expressions]);
       return '\n' + idt + this.type + (function() {
-        _a = []; _b = children;
+        _a = []; _b = this.real_children();
         for (_c = 0, _d = _b.length; _c < _d; _c++) {
           child = _b[_c];
           _a.push(child.toString(idt + TAB));
diff --git a/src/nodes.coffee b/src/nodes.coffee
index 6a813355..769dd111 100644
--- a/src/nodes.coffee
+++ b/src/nodes.coffee
@@ -94,9 +94,16 @@ BaseNode::idt: (tabs) ->
 BaseNode::contains: (block) ->
   for node in @children
     return true if block(node)
-    return true if node instanceof BaseNode and node.contains block
+    return true if node.contains and node.contains block
   false
 
+# Perform an in-order traversal of the AST.
+BaseNode::traverse: (block) ->
+  for node in @children
+    block node
+    node.traverse block if node.traverse
+
+
 # toString representation of the node, for inspecting the parse tree.
 BaseNode::toString: (idt) ->
   idt ||= ''
@@ -670,10 +677,16 @@ CodeNode: exports.CodeNode: inherit BaseNode, {
   top_sensitive: ->
     true
 
+  real_children: ->
+    flatten [@params, @body.expressions]
+
+  traverse: (block) ->
+    block this
+    block(child) for child in @real_children()
+
   toString: (idt) ->
     idt ||= ''
-    children: flatten [@params, @body.expressions]
-    '\n' + idt + @type + (child.toString(idt + TAB) for child in children).join('')
+    '\n' + idt + @type + (child.toString(idt + TAB) for child in @real_children()).join('')
 
 }