mostly implemented depth-first traversal

3 files changed, 83 insertions, 0 deletions

diff --git a/lib/arel/nodes/node.rb b/lib/arel/nodes/node.rb
index 841b954db9..90c63f0be9 100644
--- a/lib/arel/nodes/node.rb
+++ b/lib/arel/nodes/node.rb
@@ -3,6 +3,8 @@ module Arel
     ###
     # Abstract base class for all AST nodes
     class Node
+      include Enumerable
+
       ###
       # Factory method to create a Nodes::Not node that has the recipient of
       # the caller as a child.
@@ -32,6 +34,11 @@ module Arel
         viz = Visitors.for engine
         viz.accept self
       end
+
+      # Iterate through AST, nodes will be yielded depth-first
+      def each &block
+        Visitors::DepthFirst.new(block).accept self
+      end
     end
   end
 end
diff --git a/lib/arel/visitors.rb b/lib/arel/visitors.rb
index 5d4c6084b2..2b0a06d534 100644
--- a/lib/arel/visitors.rb
+++ b/lib/arel/visitors.rb
@@ -1,4 +1,5 @@
 require 'arel/visitors/visitor'
+require 'arel/visitors/depth_first'
 require 'arel/visitors/to_sql'
 require 'arel/visitors/sqlite'
 require 'arel/visitors/postgresql'
diff --git a/lib/arel/visitors/depth_first.rb b/lib/arel/visitors/depth_first.rb
new file mode 100644
index 0000000000..19796b6e72
--- /dev/null
+++ b/lib/arel/visitors/depth_first.rb
@@ -0,0 +1,75 @@
+module Arel
+  module Visitors
+    class DepthFirst < Arel::Visitors::Visitor
+      def initialize block = nil
+        @block = block || Proc.new
+      end
+
+      private
+
+      def binary o
+        visit o.left
+        visit o.right
+        @block.call o
+      end
+      alias :visit_Arel_Nodes_And                :binary
+      alias :visit_Arel_Nodes_Assignment         :binary
+      alias :visit_Arel_Nodes_Between            :binary
+      alias :visit_Arel_Nodes_DoesNotMatch       :binary
+      alias :visit_Arel_Nodes_Equality           :binary
+      alias :visit_Arel_Nodes_GreaterThan        :binary
+      alias :visit_Arel_Nodes_GreaterThanOrEqual :binary
+      alias :visit_Arel_Nodes_In                 :binary
+      alias :visit_Arel_Nodes_LessThan           :binary
+      alias :visit_Arel_Nodes_LessThanOrEqual    :binary
+      alias :visit_Arel_Nodes_Matches            :binary
+      alias :visit_Arel_Nodes_NotEqual           :binary
+      alias :visit_Arel_Nodes_NotIn              :binary
+      alias :visit_Arel_Nodes_Or                 :binary
+
+      def visit_Arel_Attribute o
+        visit o.relation
+        visit o.name
+        @block.call o
+      end
+      alias :visit_Arel_Attributes_Integer :visit_Arel_Attribute
+      alias :visit_Arel_Attributes_Float :visit_Arel_Attribute
+      alias :visit_Arel_Attributes_String :visit_Arel_Attribute
+      alias :visit_Arel_Attributes_Time :visit_Arel_Attribute
+      alias :visit_Arel_Attributes_Boolean :visit_Arel_Attribute
+      alias :visit_Arel_Attributes_Attribute :visit_Arel_Attribute
+
+      def visit_Arel_Table o
+        visit o.name
+        @block.call o
+      end
+
+      def terminal o
+        @block.call o
+      end
+      alias :visit_Arel_Nodes_SqlLiteral :terminal
+      alias :visit_Arel_SqlLiteral       :terminal
+      alias :visit_BigDecimal            :terminal
+      alias :visit_Date                  :terminal
+      alias :visit_DateTime              :terminal
+      alias :visit_FalseClass            :terminal
+      alias :visit_Fixnum                :terminal
+      alias :visit_Float                 :terminal
+      alias :visit_NilClass              :terminal
+      alias :visit_String                :terminal
+      alias :visit_Symbol                :terminal
+      alias :visit_Time                  :terminal
+      alias :visit_TrueClass             :terminal
+
+      def visit_Array o
+        o.each { |i| visit i }
+        @block.call o
+      end
+
+      def visit_Hash o
+        o.each { |k,v| visit(k); visit(v) }
+        @block.call o
+      end
+    end
+  end
+end