Refactor Preloader Code

3 files changed, 40 insertions, 81 deletions

diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index 607d376a08..b0c41c7d5f 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -17,7 +17,7 @@ module ActiveRecord
         end
 
         def run(preloader)
-          associated_records_by_owner(preloader).each do |owner, records|
+          associated_records_by_owner(preloader) do |owner, records|
             associate_records_to_owner(owner, records)
           end
         end
@@ -41,7 +41,7 @@ module ActiveRecord
             end
 
             owners.each_with_object({}) do |owner, result|
-              result[owner] = records[convert_key(owner[owner_key_name])] || []
+              yield(owner, records[convert_key(owner[owner_key_name])] || [])
             end
           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 b16fca7dc9..1bb7e98231 100644
--- a/activerecord/lib/active_record/associations/preloader/through_association.rb
+++ b/activerecord/lib/active_record/associations/preloader/through_association.rb
@@ -13,86 +13,45 @@ module ActiveRecord
         end
 
         def associated_records_by_owner(preloader)
-          already_loaded = owners.first.association(through_reflection.name).loaded?
-          through_scope = through_scope()
-
-          unless already_loaded
-            preloader.preload(owners, through_reflection.name, through_scope)
-          end
-
-          through_records = owners.map do |owner|
-            center = owner.association(through_reflection.name).target
-            [owner, Array(center)]
-          end
+          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)
 
-          if already_loaded
-            if source_type = reflection.options[:source_type]
-              through_records.map! do |owner, center|
-                center = center.select do |record|
+          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
-                [owner, center]
               end
+            else
+              owner.association(through_reflection.name).reset if through_scope
             end
-          else
-            reset_association(owners, through_reflection.name, through_scope)
-          end
-
-          middle_records = through_records.flat_map(&:last)
-
-          if preload_scope
-            reflection_scope = reflection_scope().merge(preload_scope)
-          elsif reflection.scope
-            reflection_scope = reflection_scope()
-          end
-
-          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
-              end
+            result = through_records.flat_map do |record|
+              association = record.association(source_reflection.name)
+              target = association.target
+              association.reset if preload_scope
+              target
             end
-          end.tap do
-            reset_association(middle_records, source_reflection.name, preload_scope)
+            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
+            yield(owner, result)
           end
         end
 
         private
 
-          def id_to_index_map(ids)
-            id_map = {}
-            ids.each_with_index { |id, index| id_map[id] = index }
-            id_map
-          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
 
@@ -133,6 +92,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