1 files changed, 51 insertions, 19 deletions

diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb
index de06931845..2a8530af62 100644
--- a/activerecord/lib/active_record/associations/preloader/through_association.rb
+++ b/activerecord/lib/active_record/associations/preloader/through_association.rb
@@ -2,7 +2,6 @@ module ActiveRecord
   module Associations
     class Preloader
       module ThroughAssociation #:nodoc:
-
         def through_reflection
           reflection.through_reflection
         end
@@ -11,35 +10,68 @@ module ActiveRecord
           reflection.source_reflection
         end
 
-        def associated_records_by_owner
-          through_records = through_records_by_owner
+        def associated_records_by_owner(preloader)
+          preloader.preload(owners,
+                            through_reflection.name,
+                            through_scope)
 
-          Preloader.new(through_records.values.flatten, source_reflection.name, reflection_scope).run
+          through_records = owners.map do |owner|
+            association = owner.association through_reflection.name
 
-          through_records.each do |owner, records|
-            records.map! { |r| r.send(source_reflection.name) }.flatten!
-            records.compact!
+            [owner, Array(association.reader)]
           end
-        end
 
-        private
+          reset_association owners, through_reflection.name
+
+          middle_records = through_records.map { |(_,rec)| rec }.flatten
+
+          preloaders = preloader.preload(middle_records,
+                                         source_reflection.name,
+                                         reflection_scope)
 
-        def through_records_by_owner
-          Preloader.new(owners, through_reflection.name, through_scope).run
+          @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
+
+          record_offset = {}
+          @preloaded_records.each_with_index do |record,i|
+            record_offset[record] = i
+          end
 
-          Hash[owners.map do |owner|
-            through_records = Array.wrap(owner.send(through_reflection.name))
+          through_records.each_with_object({}) { |(lhs,center),records_by_owner|
+            pl_to_middle = center.group_by { |record| middle_to_pl[record] }
 
-            # Dont cache the association - we would only be caching a subset
-            if (through_scope != through_reflection.klass.unscoped) ||
-               (reflection.options[:source_type] && through_reflection.collection?)
-              owner.association(through_reflection.name).reset
+            records_by_owner[lhs] = pl_to_middle.flat_map do |pl, middles|
+              rhs_records = middles.flat_map { |r|
+                association = r.association source_reflection.name
+
+                association.reader
+              }.compact
+
+              rhs_records.sort_by { |rhs| record_offset[rhs] }
             end
+          }
+        end
+
+        private
+
+        def reset_association(owners, association_name)
+          should_reset = (through_scope != through_reflection.klass.unscoped) ||
+             (reflection.options[:source_type] && through_reflection.collection?)
 
-            [owner, through_records]
-          end]
+          # Dont cache the association - we would only be caching a subset
+          if should_reset
+            owners.each { |owner|
+              owner.association(association_name).reset
+            }
+          end
         end
 
+
         def through_scope
           scope = through_reflection.klass.unscoped