Added #traverse and two traversal methods (depth and breadth first)

lib/mongoid/tree.rb

@@ -1,3 +1,5 @@
+require 'mongoid/tree/traversal'
+
 module Mongoid # :nodoc:
   ##
   # = Mongoid::Tree
@@ -30,6 +32,8 @@ module Mongoid # :nodoc:
   # 
   module Tree
     extend ActiveSupport::Concern
+
+    include Traversal
 
     included do
       reference_many :children, :class_name => self.name, :foreign_key => :parent_id, :inverse_of => :parent
@@ -137,7 +141,7 @@ def rearrange_children?
     private
 
     def rearrange
-      if parent_id
+      if self.parent_id
         self.parent_ids = self.class.find(self.parent_id).parent_ids + [self.parent_id]
       end
 

lib/mongoid/tree/traversal.rb

@@ -0,0 +1,85 @@
+module Mongoid # :nodoc:
+  module Tree
+    ##
+    # = Mongoid::Tree::Traversal
+    #
+    # Mongoid::Tree provides a #traverse method to walk through the tree.
+    # It supports these traversal methods:
+    #
+    # * depth_first
+    # * breadth_first
+    #
+    # == Depth First Traversal
+    #
+    # See http://en.wikipedia.org/wiki/Depth-first_search for a proper description.
+    # 
+    # Given a tree like: 
+    #
+    #   node1:
+    #    - node2:
+    #      - node3
+    #    - node4:
+    #      - node5
+    #      - node6
+    #    - node7
+    #
+    # Traversing the tree using depth first traversal would visit each node in this order:
+    #
+    #   node1, node2, node3, node4, node5, node6, node7
+    #
+    # == Breadth First Traversal
+    #
+    # See http://en.wikipedia.org/wiki/Breadth-first_search for a proper description.
+    #
+    # Given a tree like:
+    #
+    #   node1:
+    #     - node2:
+    #       - node5
+    #     - node3:
+    #       - node6
+    #       - node7
+    #     - node4
+    #
+    # Traversing the tree using breadth first traversal would visit each node in this order:
+    #
+    #   node1, node2, node3, node4, node5, node6, node7
+    #
+    module Traversal
+
+      ##
+      # Traverses the tree using the given traversal method (Default is :depth_first)
+      # and passes each document node to the block.
+      #
+      # See Mongoid::Tree::Traversal for available traversal methods.
+      #
+      # Example:
+      #
+      #   results = []
+      #   root.traverse(:depth_first) do |node|
+      #     results << node
+      #   end
+      def traverse(type = :depth_first, &block)
+        raise "No block given" unless block_given?
+        send("#{type}_traversal", &block) 
+      end      
+
+      private
+
+      def depth_first_traversal(&block)
+        block.call(self)
+        self.children.each { |c| c.send(:depth_first_traversal, &block) }
+      end
+
+      def breadth_first_traversal(&block)
+        queue = [self]
+        while queue.any? do
+          node = queue.shift
+          block.call(node)
+          queue += node.children
+        end
+      end
+
+    end
+  end
+end

spec/mongoid/tree/traversal_spec.rb

@@ -0,0 +1,92 @@
+require 'spec_helper'
+
+describe Mongoid::Tree::Traversal do
+
+  describe '#traverse' do
+
+    subject { Node.new }
+
+    it "should require a block" do
+      expect { subject.traverse }.to raise_error(/No block given/)
+    end
+
+    [:depth_first].each do |method|
+      it "should support #{method} traversal" do
+        expect { subject.traverse(method) {} }.to_not raise_error
+      end
+    end
+
+    it "should complain about unsupported traversal methods" do
+      expect { subject.traverse('non_existing') {} }.to raise_error
+    end
+
+    it "should default to depth_first traversal" do
+      subject.should_receive(:depth_first_traversal)
+      subject.traverse {}
+    end
+
+  end
+
+  describe 'depth first traversal' do
+
+    it "should traverse correctly" do
+      setup_tree <<-ENDTREE
+        node1:
+          - node2:
+            - node3
+          - node4:
+            - node5
+            - node6
+          - node7
+      ENDTREE
+
+      result = []
+      node(:node1).traverse(:depth_first) { |node| result << node }
+      result.collect { |n| n.name.to_sym }.should == [:node1, :node2, :node3, :node4, :node5, :node6, :node7]
+    end
+
+    it "should traverse correctly on merged trees" do
+
+      setup_tree <<-ENDTREE
+        - node4:
+          - node5
+          - node6:
+            - node7
+          
+        - node1:
+          - node2:
+            - node3
+      ENDTREE
+
+
+      node(:node1).children << node(:node4)
+
+
+      result = []
+      node(:node1).traverse(:depth_first) { |node| result << node }
+      result.collect { |n| n.name.to_sym }.should == [:node1, :node2, :node3, :node4, :node5, :node6, :node7]
+    end
+
+  end
+
+  describe 'breadth first traversal' do
+
+    it "should traverse correctly" do
+      tree = setup_tree <<-ENDTREE
+        node1:
+          - node2:
+            - node5
+          - node3:
+            - node6
+            - node7
+          - node4
+      ENDTREE
+
+      result = []
+      node(:node1).traverse(:breadth_first) { |n| result << n }
+      result.collect { |n| n.name.to_sym }.should == [:node1, :node2, :node3, :node4, :node5, :node6, :node7]
+    end
+
+  end
+
+end