store join parts in a tree

1 files changed, 45 insertions, 16 deletions

diff --git a/activerecord/lib/active_record/associations/join_dependency.rb b/activerecord/lib/active_record/associations/join_dependency.rb
index f12d5c49b5..ce6cd1461a 100644
--- a/activerecord/lib/active_record/associations/join_dependency.rb
+++ b/activerecord/lib/active_record/associations/join_dependency.rb
@@ -4,7 +4,7 @@ module ActiveRecord
       autoload :JoinBase,        'active_record/associations/join_dependency/join_base'
       autoload :JoinAssociation, 'active_record/associations/join_dependency/join_association'
 
-      attr_reader :join_parts, :alias_tracker, :base_klass
+      attr_reader :alias_tracker, :base_klass
 
       def self.make_tree(associations)
         hash = {}
@@ -54,23 +54,48 @@ module ActiveRecord
       def initialize(base, associations, joins)
         @base_klass    = base
         @table_joins   = joins
-        @join_parts    = [JoinBase.new(base)]
+        @join_parts    = Node.new JoinBase.new(base)
         @alias_tracker = AliasTracker.new(base.connection, joins)
         @alias_tracker.aliased_name_for(base.table_name) # Updates the count for base.table_name to 1
         tree = self.class.make_tree associations
-        build tree, join_parts.last, Arel::InnerJoin
+        build tree, @join_parts, Arel::InnerJoin
+      end
+
+      class Node # :nodoc:
+        include Enumerable
+
+        attr_reader :join_part, :children
+
+        def initialize(join_part)
+          @join_part = join_part
+          @children  = []
+        end
+
+        def each
+          yield self
+          iter = lambda { |list|
+            list.each { |item|
+              yield item
+              iter.call item.children
+            }
+          }
+          iter.call children
+        end
+      end
+
+      def join_parts
+        @join_parts.map(&:join_part)
       end
 
       def graft(*associations)
         join_assocs = join_associations
-        base        = join_base
 
         associations.reject { |association|
           join_assocs.detect { |a| association == a }
         }.each { |association|
-          join_part = find_parent_part(association.parent) || base
+          join_node = find_parent_node(association.parent) || @join_part
           type      = association.join_type
-          find_or_build_scalar association.reflection, join_part, type
+          find_or_build_scalar association.reflection, join_node, type
         }
         self
       end
@@ -124,19 +149,20 @@ module ActiveRecord
         end
       end
 
-      def find_parent_part(parent)
-        join_parts.detect do |join_part|
+      def find_parent_node(parent)
+        @join_parts.find { |node|
+          join_part = node.join_part
           case parent
           when JoinBase
             parent.base_klass == join_part.base_klass
           else
             parent == join_part
           end
-        end
+        }
       end
 
       def join_base
-        join_parts.first
+        @join_parts.first.join_part
       end
 
       def remove_duplicate_results!(base, records, associations)
@@ -179,19 +205,21 @@ module ActiveRecord
           raise ConfigurationError, "Association named '#{ name }' was not found on #{ klass.name }; perhaps you misspelled it?"
       end
 
-      def build(associations, parent, join_type)
+      def build(associations, root, join_type)
+        parent = root.join_part
         associations.each do |name, right|
           reflection = find_reflection parent.base_klass, name
           join_association = build_join_association reflection, parent, join_type
-          @join_parts << join_association
+          root.children << join_association
           build right, join_association, join_type
         end
       end
 
-      def find_or_build_scalar(reflection, parent, join_type)
-        unless join_association = find_join_association(reflection, parent)
+      def find_or_build_scalar(reflection, node, join_type)
+        parent = node.join_part
+        unless join_association = find_join_association(reflection, node.join_part)
           join_association = build_join_association(reflection, parent, join_type)
-          @join_parts << join_association
+          node.children << join_association
         end
         join_association
       end
@@ -215,7 +243,8 @@ module ActiveRecord
           raise EagerLoadPolymorphicError.new(reflection)
         end
 
-        JoinAssociation.new(reflection, join_parts.length, parent, join_type, alias_tracker)
+        part = JoinAssociation.new(reflection, join_parts.length, parent, join_type, alias_tracker)
+        Node.new part
       end
 
       def construct(parent, associations, join_parts, row, rs)