mostly implemented depth-first traversal

4 files changed, 159 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
diff --git a/test/visitors/test_depth_first.rb b/test/visitors/test_depth_first.rb
new file mode 100644
index 0000000000..eff60576fc
--- /dev/null
+++ b/test/visitors/test_depth_first.rb
@@ -0,0 +1,76 @@
+require 'helper'
+
+module Arel
+  module Visitors
+    class TestDepthFirst < MiniTest::Unit::TestCase
+      Collector = Struct.new(:calls) do
+        def call object
+          calls << object
+        end
+      end
+
+      def setup
+        @collector = Collector.new []
+        @visitor = Visitors::DepthFirst.new @collector
+      end
+
+      [
+        Arel::Nodes::And,
+        Arel::Nodes::Assignment,
+        Arel::Nodes::Between,
+        Arel::Nodes::DoesNotMatch,
+        Arel::Nodes::Equality,
+        Arel::Nodes::GreaterThan,
+        Arel::Nodes::GreaterThanOrEqual,
+        Arel::Nodes::In,
+        Arel::Nodes::LessThan,
+        Arel::Nodes::LessThanOrEqual,
+        Arel::Nodes::Matches,
+        Arel::Nodes::NotEqual,
+        Arel::Nodes::NotIn,
+        Arel::Nodes::Or,
+      ].each do |klass|
+        define_method("test_#{klass.name.gsub('::', '_')}") do
+          binary = klass.new(:a, :b)
+          @visitor.accept binary
+          assert_equal [:a, :b, binary], @collector.calls
+        end
+      end
+
+      [
+        Arel::Attributes::Integer,
+        Arel::Attributes::Float,
+        Arel::Attributes::String,
+        Arel::Attributes::Time,
+        Arel::Attributes::Boolean,
+        Arel::Attributes::Attribute
+      ].each do |klass|
+        define_method("test_#{klass.name.gsub('::', '_')}") do
+          binary = klass.new(:a, :b)
+          @visitor.accept binary
+          assert_equal [:a, :b, binary], @collector.calls
+        end
+      end
+
+      def test_table
+        relation = Arel::Table.new(:users)
+        @visitor.accept relation
+        assert_equal ['users', relation], @collector.calls
+      end
+
+      def test_array
+        node = Nodes::Or.new(:a, :b)
+        list = [node]
+        @visitor.accept list
+        assert_equal [:a, :b, node, list], @collector.calls
+      end
+
+      def test_hash
+        node = Nodes::Or.new(:a, :b)
+        hash = { node => node }
+        @visitor.accept hash
+        assert_equal [:a, :b, node, :a, :b, node, hash], @collector.calls
+      end
+    end
+  end
+end