keep preloaded records in a list rather than extract from a hash

3 files changed, 21 insertions, 16 deletions

diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index cc8eefb234..bd63b40472 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -3,6 +3,7 @@ module ActiveRecord
     class Preloader
       class Association #:nodoc:
         attr_reader :owners, :reflection, :preload_scope, :model, :klass
+        attr_reader :preloaded_records
 
         def initialize(klass, owners, reflection, preload_scope)
           @klass         = klass
@@ -12,7 +13,7 @@ module ActiveRecord
           @model         = owners.first && owners.first.class
           @scope         = nil
           @owners_by_key = nil
-          @associated_records_by_owner = nil
+          @preloaded_records = []
         end
 
         def run(preloader)
@@ -65,15 +66,9 @@ module ActiveRecord
           reflection.options
         end
 
-        def preloaded_records
-          @associated_records_by_owner.values.flatten
-        end
-
         private
 
         def associated_records_by_owner(preloader)
-          return @associated_records_by_owner if @associated_records_by_owner
-
           owners_map = owners_by_key
           owner_keys = owners_map.keys.compact
 
@@ -94,18 +89,19 @@ module ActiveRecord
               end
             end
           end
+
           owners.each_with_object(records_by_owner) do |owner,h|
             h[owner] ||= []
           end
-
-          @associated_records_by_owner = records_by_owner
         end
 
         def load_slices(slices)
-          slices.flat_map { |slice|
-            records_for(slice).to_a.map! { |record|
-              [record, record[association_key_name]]
-            }
+          @preloaded_records = slices.flat_map { |slice|
+            records_for(slice)
+          }
+
+          @preloaded_records.map { |record|
+            [record, record[association_key_name]]
           }
         end
 
diff --git a/activerecord/lib/active_record/associations/preloader/has_and_belongs_to_many.rb b/activerecord/lib/active_record/associations/preloader/has_and_belongs_to_many.rb
index 996575f5aa..b62ca6f681 100644
--- a/activerecord/lib/active_record/associations/preloader/has_and_belongs_to_many.rb
+++ b/activerecord/lib/active_record/associations/preloader/has_and_belongs_to_many.rb
@@ -38,7 +38,7 @@ module ActiveRecord
           caster = nil
           name = association_key_name
 
-          slices.flat_map { |slice|
+          records_to_keys = slices.flat_map { |slice|
             records = records_for(slice)
             caster ||= records.column_types.fetch(name, records.identity_type)
             records.map! { |row|
@@ -46,6 +46,9 @@ module ActiveRecord
               [record, caster.type_cast(row[name])]
             }
           }
+          @preloaded_records = records_to_keys.map(&:first)
+
+          records_to_keys
         end
 
         def build_scope
diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb
index 7187cf38fd..ec67a52885 100644
--- a/activerecord/lib/active_record/associations/preloader/through_association.rb
+++ b/activerecord/lib/active_record/associations/preloader/through_association.rb
@@ -2,6 +2,10 @@ module ActiveRecord
   module Associations
     class Preloader
       module ThroughAssociation #:nodoc:
+        def initialize(klass, owners, reflection, preload_scope)
+          super
+          @associated_records_by_owner = nil
+        end
 
         def through_reflection
           reflection.through_reflection
@@ -11,9 +15,11 @@ module ActiveRecord
           reflection.source_reflection
         end
 
-        def associated_records_by_owner(preloader)
-          return @associated_records_by_owner if @associated_records_by_owner
+        def preloaded_records
+          @associated_records_by_owner.values.flatten
+        end
 
+        def associated_records_by_owner(preloader)
           preloader.preload(owners,
                             through_reflection.name,
                             through_scope)