From 56b779e19deeddf59279217e9596210b232b730a Mon Sep 17 00:00:00 2001
From: benedikt <benedikt@synatic.net>
Date: Mon, 26 Jul 2010 17:14:14 +0200
Subject: [PATCH] Added #traverse and two traversal methods (depth and breadth
 first)

---
 lib/mongoid/tree.rb                 |  6 +-
 lib/mongoid/tree/traversal.rb       | 85 ++++++++++++++++++++++++++
 spec/mongoid/tree/traversal_spec.rb | 92 +++++++++++++++++++++++++++++
 3 files changed, 182 insertions(+), 1 deletion(-)
 create mode 100644 lib/mongoid/tree/traversal.rb
 create mode 100644 spec/mongoid/tree/traversal_spec.rb

diff --git a/lib/mongoid/tree.rb b/lib/mongoid/tree.rb
index b6563e1..4041620 100644
--- a/lib/mongoid/tree.rb
+++ b/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
     
diff --git a/lib/mongoid/tree/traversal.rb b/lib/mongoid/tree/traversal.rb
new file mode 100644
index 0000000..4e510e0
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/spec/mongoid/tree/traversal_spec.rb b/spec/mongoid/tree/traversal_spec.rb
new file mode 100644
index 0000000..e093433
--- /dev/null
+++ b/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
\ No newline at end of file