An extension of mongoid-tree gem that I used in implementing Decision Tree for the Foreplot.com. Optimisation to speed up access to and execution of block on each leaf nodes in a tree. Results in orders of magnitude in speed improvement.

tree_ext.rb

# Copyright (c) 2012 Medhi Widjaja

## This is an extension of Mongoid::Tree module in mongoid-tree gem

module Mongoid
  module TreeExt
    
    ##
    # Returns an Array containing all leaves of this document. It's cheaper than leaves method.
    # Caution: Can't be chained, since it returns Array instead of Mongoid Criteria object

    def leaf_nodes
      result = [self] if self.leaf?
      if result.nil?
        result =  self.children.collect {|child| child.leaf_nodes }
      else
        result << self.children.collect {|child| child.leaf_nodes }
      end
      result.flatten
    end

    ##
    # Executes block on the leaf nodes under this document. This is much
    # faster (hundreds times faster) than doing node.leaves.each {}
    # It's also 2 times faster than doing node.leaf_nodes.each {}

    def each_leaf(&block)
      block.call(self) if self.leaf?
      self.children.each {|child| child.each_leaf(&block)}
    end

  end
end