Merge pull request #18279 from eileencodes/refactor-association-scope

Refactoring of add_constraints in AssociationScope

4 files changed, 169 insertions, 84 deletions

diff --git a/activerecord/lib/active_record/associations/alias_tracker.rb b/activerecord/lib/active_record/associations/alias_tracker.rb
index f9c9f8afda..2b7e4f28c5 100644
--- a/activerecord/lib/active_record/associations/alias_tracker.rb
+++ b/activerecord/lib/active_record/associations/alias_tracker.rb
@@ -5,20 +5,23 @@ module ActiveRecord
     # Keeps track of table aliases for ActiveRecord::Associations::ClassMethods::JoinDependency and
     # ActiveRecord::Associations::ThroughAssociationScope
     class AliasTracker # :nodoc:
-      attr_reader :aliases, :connection
+      attr_reader :aliases
 
-      def self.empty(connection)
-        new connection, Hash.new(0)
+      def self.create(connection, initial_table, type_caster)
+        aliases = Hash.new(0)
+        aliases[initial_table] = 1
+        new connection, aliases, type_caster
       end
 
-      def self.create(connection, table_joins)
-        if table_joins.empty?
-          empty connection
+      def self.create_with_joins(connection, initial_table, joins, type_caster)
+        if joins.empty?
+          create(connection, initial_table, type_caster)
         else
-          aliases = Hash.new { |h,k|
-            h[k] = initial_count_for(connection, k, table_joins)
+          aliases = Hash.new { |h, k|
+            h[k] = initial_count_for(connection, k, joins)
           }
-          new connection, aliases
+          aliases[initial_table] = 1
+          new connection, aliases, type_caster
         end
       end
 
@@ -51,19 +54,20 @@ module ActiveRecord
       end
 
       # table_joins is an array of arel joins which might conflict with the aliases we assign here
-      def initialize(connection, aliases)
+      def initialize(connection, aliases, type_caster)
         @aliases    = aliases
         @connection = connection
+        @type_caster = type_caster
       end
 
-      def aliased_table_for(table_name, aliased_name, **table_options)
+      def aliased_table_for(table_name, aliased_name)
         if aliases[table_name].zero?
           # If it's zero, we can have our table_name
           aliases[table_name] = 1
-          Arel::Table.new(table_name, table_options)
+          Arel::Table.new(table_name, type_caster: @type_caster)
         else
           # Otherwise, we need to use an alias
-          aliased_name = connection.table_alias_for(aliased_name)
+          aliased_name = @connection.table_alias_for(aliased_name)
 
           # Update the count
           aliases[aliased_name] += 1
@@ -73,14 +77,14 @@ module ActiveRecord
           else
             aliased_name
           end
-          Arel::Table.new(table_name, table_options).alias(table_alias)
+          Arel::Table.new(table_name, type_caster: @type_caster).alias(table_alias)
         end
       end
 
       private
 
         def truncate(name)
-          name.slice(0, connection.table_alias_length - 2)
+          name.slice(0, @connection.table_alias_length - 2)
         end
     end
   end
diff --git a/activerecord/lib/active_record/associations/association_scope.rb b/activerecord/lib/active_record/associations/association_scope.rb
index 53f65920e1..d06b7b3508 100644
--- a/activerecord/lib/active_record/associations/association_scope.rb
+++ b/activerecord/lib/active_record/associations/association_scope.rb
@@ -33,10 +33,11 @@ module ActiveRecord
         reflection    = association.reflection
         scope         = klass.unscoped
         owner         = association.owner
-        alias_tracker = AliasTracker.empty connection
+        alias_tracker = AliasTracker.create connection, association.klass.table_name, klass.type_caster
+        chain_head, chain_tail = get_chain(reflection, association, alias_tracker)
 
         scope.extending! Array(reflection.options[:extend])
-        add_constraints(scope, owner, klass, reflection, alias_tracker)
+        add_constraints(scope, owner, klass, reflection, connection, chain_head, chain_tail)
       end
 
       def join_type
@@ -61,23 +62,6 @@ module ActiveRecord
       end
 
       private
-
-      def construct_tables(chain, klass, refl, alias_tracker)
-        chain.map do |reflection|
-          alias_tracker.aliased_table_for(
-            table_name_for(reflection, klass, refl),
-            table_alias_for(reflection, refl, reflection != refl),
-            type_caster: klass.type_caster,
-          )
-        end
-      end
-
-      def table_alias_for(reflection, refl, join = false)
-        name = "#{reflection.plural_name}_#{alias_suffix(refl)}"
-        name << "_join" if join
-        name
-      end
-
       def join(table, constraint)
         table.create_join(table, table.create_on(constraint), join_type)
       end
@@ -96,8 +80,8 @@ module ActiveRecord
         bind_value scope, column, value, connection
       end
 
-      def last_chain_scope(scope, table, reflection, owner, connection, assoc_klass)
-        join_keys = reflection.join_keys(assoc_klass)
+      def last_chain_scope(scope, table, reflection, owner, connection, association_klass)
+        join_keys = reflection.join_keys(association_klass)
         key = join_keys.key
         foreign_key = join_keys.foreign_key
 
@@ -113,8 +97,8 @@ module ActiveRecord
         end
       end
 
-      def next_chain_scope(scope, table, reflection, connection, assoc_klass, foreign_table, next_reflection)
-        join_keys = reflection.join_keys(assoc_klass)
+      def next_chain_scope(scope, table, reflection, connection, association_klass, foreign_table, next_reflection)
+        join_keys = reflection.join_keys(association_klass)
         key = join_keys.key
         foreign_key = join_keys.foreign_key
 
@@ -129,38 +113,57 @@ module ActiveRecord
         scope = scope.joins(join(foreign_table, constraint))
       end
 
-      def add_constraints(scope, owner, assoc_klass, refl, tracker)
-        chain = refl.chain
-        scope_chain = refl.scope_chain
-        connection = tracker.connection
+      class ReflectionProxy < SimpleDelegator # :nodoc:
+        attr_accessor :next
+        attr_reader :alias_name
+
+        def initialize(reflection, alias_name)
+          super(reflection)
+          @alias_name = alias_name
+        end
+
+        def all_includes; nil; end
+      end
 
-        tables = construct_tables(chain, assoc_klass, refl, tracker)
+      def get_chain(reflection, association, tracker)
+        name = reflection.name
+        runtime_reflection = Reflection::RuntimeReflection.new(reflection, association)
+        previous_reflection = runtime_reflection
+        reflection.chain.drop(1).each do |refl|
+          alias_name = tracker.aliased_table_for(refl.table_name, refl.alias_candidate(name))
+          proxy = ReflectionProxy.new(refl, alias_name)
+          previous_reflection.next = proxy
+          previous_reflection = proxy
+        end
+        [runtime_reflection, previous_reflection]
+      end
 
-        owner_reflection = chain.last
-        table = tables.last
-        scope = last_chain_scope(scope, table, owner_reflection, owner, connection, assoc_klass)
+      def add_constraints(scope, owner, association_klass, refl, connection, chain_head, chain_tail)
+        owner_reflection = chain_tail
+        table = owner_reflection.alias_name
+        scope = last_chain_scope(scope, table, owner_reflection, owner, connection, association_klass)
 
-        chain.each_with_index do |reflection, i|
-          table, foreign_table = tables.shift, tables.first
+        reflection = chain_head
+        loop do
+          break unless reflection
+          table = reflection.alias_name
 
-          unless reflection == chain.last
-            next_reflection = chain[i + 1]
-            scope = next_chain_scope(scope, table, reflection, connection, assoc_klass, foreign_table, next_reflection)
+          unless reflection == chain_tail
+            next_reflection = reflection.next
+            foreign_table = next_reflection.alias_name
+            scope = next_chain_scope(scope, table, reflection, connection, association_klass, foreign_table, next_reflection)
           end
 
-          is_first_chain = i == 0
-          klass = is_first_chain ? assoc_klass : reflection.klass
-
           # Exclude the scope of the association itself, because that
           # was already merged in the #scope method.
-          scope_chain[i].each do |scope_chain_item|
-            item  = eval_scope(klass, scope_chain_item, owner)
+          reflection.constraints.each do |scope_chain_item|
+            item  = eval_scope(reflection.klass, scope_chain_item, owner)
 
             if scope_chain_item == refl.scope
               scope.merge! item.except(:where, :includes, :bind)
             end
 
-            if is_first_chain
+            reflection.all_includes do
               scope.includes! item.includes_values
             end
 
@@ -168,26 +171,13 @@ module ActiveRecord
             scope.bind_values  += item.bind_values
             scope.order_values |= item.order_values
           end
+
+          reflection = reflection.next
         end
 
         scope
       end
 
-      def alias_suffix(refl)
-        refl.name
-      end
-
-      def table_name_for(reflection, klass, refl)
-        if reflection == refl
-          # If this is a polymorphic belongs_to, we want to get the klass from the
-          # association because it depends on the polymorphic_type attribute of
-          # the owner
-          klass.table_name
-        else
-          reflection.table_name
-        end
-      end
-
       def eval_scope(klass, scope, owner)
         klass.unscoped.instance_exec(owner, &scope)
       end
diff --git a/activerecord/lib/active_record/associations/join_dependency.rb b/activerecord/lib/active_record/associations/join_dependency.rb
index 66e997c3c8..4b75370171 100644
--- a/activerecord/lib/active_record/associations/join_dependency.rb
+++ b/activerecord/lib/active_record/associations/join_dependency.rb
@@ -93,8 +93,7 @@ module ActiveRecord
       #    joins # =>  []
       #
       def initialize(base, associations, joins)
-        @alias_tracker = AliasTracker.create(base.connection, joins)
-        @alias_tracker.aliased_table_for(base.table_name, base.table_name, type_caster: base.type_caster) # Updates the count for base.table_name to 1
+        @alias_tracker = AliasTracker.create_with_joins(base.connection, base.table_name, joins, base.type_caster)
         tree = self.class.make_tree associations
         @join_root = JoinBase.new base, build(tree, base)
         @join_root.children.each { |child| construct_tables! @join_root, child }
@@ -186,13 +185,9 @@ module ActiveRecord
 
       def table_aliases_for(parent, node)
         node.reflection.chain.map { |reflection|
-          if reflection.klass
-            type_caster = reflection.klass.type_caster
-          end
           alias_tracker.aliased_table_for(
             reflection.table_name,
-            table_alias_for(reflection, parent, reflection != node.reflection),
-            type_caster: type_caster,
+            table_alias_for(reflection, parent, reflection != node.reflection)
           )
         }
       end
diff --git a/activerecord/lib/active_record/reflection.rb b/activerecord/lib/active_record/reflection.rb
index dd746a4e10..7696ef13c7 100644
--- a/activerecord/lib/active_record/reflection.rb
+++ b/activerecord/lib/active_record/reflection.rb
@@ -149,7 +149,7 @@ module ActiveRecord
 
       JoinKeys = Struct.new(:key, :foreign_key) # :nodoc:
 
-      def join_keys(assoc_klass)
+      def join_keys(association_klass)
         JoinKeys.new(foreign_key, active_record_primary_key)
       end
 
@@ -161,7 +161,16 @@ module ActiveRecord
 
         macro
       end
+
+      def constraints
+        scope_chain.flatten
+      end
+
+      def alias_candidate(name)
+        "#{plural_name}_#{name}"
+      end
     end
+
     # Base class for AggregateReflection and AssociationReflection. Objects of
     # AggregateReflection and AssociationReflection are returned by the Reflection::ClassMethods.
     #
@@ -601,8 +610,8 @@ module ActiveRecord
 
       def belongs_to?; true; end
 
-      def join_keys(assoc_klass)
-        key = polymorphic? ? association_primary_key(assoc_klass) : association_primary_key
+      def join_keys(association_klass)
+        key = polymorphic? ? association_primary_key(association_klass) : association_primary_key
         JoinKeys.new(key, foreign_key)
       end
 
@@ -698,6 +707,11 @@ module ActiveRecord
         @chain ||= begin
           a = source_reflection.chain
           b = through_reflection.chain
+
+          if options[:source_type]
+            b[0] = PolymorphicReflection.new(b[0], self)
+          end
+
           chain = a + b
           chain[0] = self # Use self so we don't lose the information from :source_type
           chain
@@ -745,8 +759,8 @@ module ActiveRecord
         end
       end
 
-      def join_keys(assoc_klass)
-        source_reflection.join_keys(assoc_klass)
+      def join_keys(association_klass)
+        source_reflection.join_keys(association_klass)
       end
 
       # The macro used by the source association
@@ -855,6 +869,12 @@ module ActiveRecord
         check_validity_of_inverse!
       end
 
+      def constraints
+        scope_chain = source_reflection.constraints
+        scope_chain << scope if scope
+        scope_chain
+      end
+
       protected
 
         def actual_source_reflection # FIXME: this is a horrible name
@@ -877,5 +897,81 @@ module ActiveRecord
         delegate(*delegate_methods, to: :delegate_reflection)
 
     end
+
+    class PolymorphicReflection < ThroughReflection # :nodoc:
+      def initialize(reflection, previous_reflection)
+        @reflection = reflection
+        @previous_reflection = previous_reflection
+      end
+
+      def klass
+        @reflection.klass
+      end
+
+      def scope
+        @reflection.scope
+      end
+
+      def table_name
+        @reflection.table_name
+      end
+
+      def plural_name
+        @reflection.plural_name
+      end
+
+      def join_keys(association_klass)
+        @reflection.join_keys(association_klass)
+      end
+
+      def type
+        @reflection.type
+      end
+
+      def constraints
+        [source_type_info]
+      end
+
+      def source_type_info
+        type = @previous_reflection.foreign_type
+        source_type = @previous_reflection.options[:source_type]
+        lambda { |object| where(type => source_type) }
+      end
+    end
+
+    class RuntimeReflection < PolymorphicReflection # :nodoc:
+      attr_accessor :next
+
+      def initialize(reflection, association)
+        @reflection = reflection
+        @association = association
+      end
+
+      def klass
+        @association.klass
+      end
+
+      def table_name
+        klass.table_name
+      end
+
+      def constraints
+        @reflection.constraints
+      end
+
+      def source_type_info
+        @reflection.source_type_info
+      end
+
+      def alias_candidate(name)
+        "#{plural_name}_#{name}_join"
+      end
+
+      def alias_name
+        Arel::Table.new(table_name)
+      end
+
+      def all_includes; yield; end
+    end
   end
 end