Merge branch 'master' into activestorage-guide

4 files changed, 60 insertions, 87 deletions

diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb
index e1754d4a19..5a93a89d0a 100644
--- a/activerecord/lib/active_record/associations/preloader.rb
+++ b/activerecord/lib/active_record/associations/preloader.rb
@@ -166,8 +166,6 @@ module ActiveRecord
         end
 
         class AlreadyLoaded # :nodoc:
-          attr_reader :owners, :reflection
-
           def initialize(klass, owners, reflection, preload_scope)
             @owners = owners
             @reflection = reflection
@@ -178,6 +176,9 @@ module ActiveRecord
           def preloaded_records
             owners.flat_map { |owner| owner.association(reflection.name).target }
           end
+
+          protected
+            attr_reader :owners, :reflection
         end
 
         # Returns a class containing the logic needed to load preload the data
diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index 607d376a08..19c337dc39 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -4,7 +4,6 @@ module ActiveRecord
   module Associations
     class Preloader
       class Association #:nodoc:
-        attr_reader :owners, :reflection, :preload_scope, :model, :klass
         attr_reader :preloaded_records
 
         def initialize(klass, owners, reflection, preload_scope)
@@ -17,11 +16,20 @@ module ActiveRecord
         end
 
         def run(preloader)
-          associated_records_by_owner(preloader).each do |owner, records|
-            associate_records_to_owner(owner, records)
+          records = load_records do |record|
+            owner = owners_by_key[convert_key(record[association_key_name])]
+            association = owner.association(reflection.name)
+            association.set_inverse_instance(record)
+          end
+
+          owners.each do |owner|
+            associate_records_to_owner(owner, records[convert_key(owner[owner_key_name])] || [])
           end
         end
 
+        protected
+          attr_reader :owners, :reflection, :preload_scope, :model, :klass
+
         private
           # The name of the key on the associated records
           def association_key_name
@@ -33,18 +41,6 @@ module ActiveRecord
             reflection.join_foreign_key
           end
 
-          def associated_records_by_owner(preloader)
-            records = load_records do |record|
-              owner = owners_by_key[convert_key(record[association_key_name])]
-              association = owner.association(reflection.name)
-              association.set_inverse_instance(record)
-            end
-
-            owners.each_with_object({}) do |owner, result|
-              result[owner] = records[convert_key(owner[owner_key_name])] || []
-            end
-          end
-
           def associate_records_to_owner(owner, records)
             raise NotImplementedError
           end
diff --git a/activerecord/lib/active_record/associations/preloader/has_many_through.rb b/activerecord/lib/active_record/associations/preloader/has_many_through.rb
index 0639fdca44..3e17d07a33 100644
--- a/activerecord/lib/active_record/associations/preloader/has_many_through.rb
+++ b/activerecord/lib/active_record/associations/preloader/has_many_through.rb
@@ -5,16 +5,6 @@ module ActiveRecord
     class Preloader
       class HasManyThrough < CollectionAssociation #:nodoc:
         include ThroughAssociation
-
-        def associated_records_by_owner(preloader)
-          records_by_owner = super
-
-          if reflection_scope.distinct_value
-            records_by_owner.each_value(&:uniq!)
-          end
-
-          records_by_owner
-        end
       end
     end
   end
diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb
index fa32cc5553..762275fbad 100644
--- a/activerecord/lib/active_record/associations/preloader/through_association.rb
+++ b/activerecord/lib/active_record/associations/preloader/through_association.rb
@@ -4,77 +4,53 @@ module ActiveRecord
   module Associations
     class Preloader
       module ThroughAssociation #:nodoc:
-        def through_reflection
-          reflection.through_reflection
-        end
-
-        def source_reflection
-          reflection.source_reflection
-        end
-
-        def associated_records_by_owner(preloader)
-          through_scope = through_scope()
-
-          preloader.preload(owners,
-                            through_reflection.name,
-                            through_scope)
-
-          through_records = owners.map do |owner|
-            center = owner.association(through_reflection.name).target
-            [owner, Array(center)]
-          end
-
-          reset_association(owners, through_reflection.name, through_scope)
-
-          middle_records = through_records.flat_map(&:last)
-
-          reflection_scope = reflection_scope() if reflection.scope
-
-          preloaders = preloader.preload(middle_records,
-                                         source_reflection.name,
-                                         reflection_scope)
-
+        def run(preloader)
+          already_loaded     = owners.first.association(through_reflection.name).loaded?
+          through_scope      = through_scope()
+          reflection_scope   = target_reflection_scope
+          through_preloaders = preloader.preload(owners, through_reflection.name, through_scope)
+          middle_records     = through_preloaders.flat_map(&:preloaded_records)
+          preloaders         = preloader.preload(middle_records, source_reflection.name, reflection_scope)
           @preloaded_records = preloaders.flat_map(&:preloaded_records)
 
-          middle_to_pl = preloaders.each_with_object({}) do |pl, h|
-            pl.owners.each { |middle|
-              h[middle] = pl
-            }
-          end
-
-          through_records.each_with_object({}) do |(lhs, center), records_by_owner|
-            pl_to_middle = center.group_by { |record| middle_to_pl[record] }
-
-            records_by_owner[lhs] = pl_to_middle.flat_map do |pl, middles|
-              rhs_records = middles.flat_map { |r|
-                r.association(source_reflection.name).target
-              }.compact
-
-              # Respect the order on `reflection_scope` if it exists, else use the natural order.
-              if reflection_scope && !reflection_scope.order_values.empty?
-                @id_map ||= id_to_index_map @preloaded_records
-                rhs_records.sort_by { |rhs| @id_map[rhs] }
-              else
-                rhs_records
+          owners.each do |owner|
+            through_records = Array(owner.association(through_reflection.name).target)
+            if already_loaded
+              if source_type = reflection.options[:source_type]
+                through_records = through_records.select do |record|
+                  record[reflection.foreign_type] == source_type
+                end
               end
+            else
+              owner.association(through_reflection.name).reset if through_scope
+            end
+            result = through_records.flat_map do |record|
+              association = record.association(source_reflection.name)
+              target = association.target
+              association.reset if preload_scope
+              target
             end
+            result.compact!
+            if reflection_scope
+              result.sort_by! { |rhs| preload_index[rhs] } if reflection_scope.order_values.any?
+              result.uniq! if reflection_scope.distinct_value
+            end
+            associate_records_to_owner(owner, result)
           end
         end
 
         private
+          def through_reflection
+            reflection.through_reflection
+          end
 
-          def id_to_index_map(ids)
-            id_map = {}
-            ids.each_with_index { |id, index| id_map[id] = index }
-            id_map
+          def source_reflection
+            reflection.source_reflection
           end
 
-          def reset_association(owners, association_name, should_reset)
-            # Don't cache the association - we would only be caching a subset
-            if should_reset
-              owners.each { |owner|
-                owner.association(association_name).reset
-              }
+          def preload_index
+            @preload_index ||= @preloaded_records.each_with_object({}).with_index do |(id, result), index|
+              result[id] = index
             end
           end
 
@@ -115,6 +91,16 @@ module ActiveRecord
 
             scope unless scope.empty_scope?
           end
+
+          def target_reflection_scope
+            if preload_scope
+              reflection_scope.merge(preload_scope)
+            elsif reflection.scope
+              reflection_scope
+            else
+              nil
+            end
+          end
       end
     end
   end