push `run` up to preload

3 files changed, 41 insertions, 14 deletions

diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb
index c7ab926f56..03c66cab4f 100644
--- a/activerecord/lib/active_record/associations/preloader.rb
+++ b/activerecord/lib/active_record/associations/preloader.rb
@@ -105,22 +105,36 @@ module ActiveRecord
       def preload(association, records)
         case association
         when Hash
-          preload_hash(association, records)
+          preload_hash(association, records).each(&:run)
         when Symbol
-          preload_one(association, records)
+          preloaders_for_one(association, records).each(&:run)
         when String
-          preload_one(association.to_sym, records)
+          preloaders_for_one(association.to_sym, records).each(&:run)
         else
           raise ArgumentError, "#{association.inspect} was not recognised for preload"
         end
       end
 
-      def preload_hash(association, records)
+      def preloaders_for_hash(association, records)
         parent, child = association.to_a.first # hash should only be of length 1
 
-        preload_one parent, records
-        run_preload Array.wrap(child),
-          records.map { |record| record.send(parent) }.flatten.compact.uniq
+        loaders = preloaders_for_one parent, records
+
+        recs = loaders.flat_map(&:target_records).uniq
+        lls = Array.wrap(child).flat_map { |assoc|
+          case assoc
+          when Hash then preloaders_for_hash(assoc, recs)
+          when Symbol then preloaders_for_one(assoc, recs)
+          when String then preloaders_for_one(assoc.to_sym, recs)
+          else
+            raise
+          end
+        }
+        loaders + lls
+      end
+
+      def preload_hash(association, records)
+        preloaders_for_hash(association, records).each(&:run)
       end
 
       # Not all records have the same class, so group then preload group on the reflection
@@ -131,9 +145,15 @@ module ActiveRecord
       # classes, depending on the polymorphic_type field. So we group by the classes as
       # well.
       def preload_one(association, records)
-        grouped_records(association, records).each do |reflection, klasses|
-          klasses.each do |klass, rs|
-            preloader_for(reflection).new(klass, rs, reflection, preload_scope).run
+        preloaders_for_one(association, records).each { |loader|
+          loader.run
+        }
+      end
+
+      def preloaders_for_one(association, records)
+        grouped_records(association, records).flat_map do |reflection, klasses|
+          klasses.map do |rhs_klass, rs|
+            preloader_for(reflection).new(rhs_klass, rs, reflection, preload_scope)
           end
         end
       end
diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index 051bc3a8da..9b8750b9ea 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -72,6 +72,12 @@ module ActiveRecord
         end
 
         def target_records
+          associated_records_by_owner.values.flatten
+        end
+
+        private
+
+        def lhs_records
           return @target_records if @target_records
 
           owners_map = owners_by_key
@@ -85,8 +91,6 @@ module ActiveRecord
           }
         end
 
-        private
-
         def associated_records_by_owner
           owners_map = owners_by_key
           owner_keys = owners_map.keys.compact
@@ -98,7 +102,7 @@ module ActiveRecord
             # Some databases impose a limit on the number of ids in a list (in Oracle it's 1000)
             # Make several smaller queries if necessary or make one query if the adapter supports it
             caster = type_caster
-            target_records.each do |record|
+            lhs_records.each do |record|
               owner_key = caster.call record[association_key_name]
 
               owners_map[owner_key].each do |owner|
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 87992a70e8..48ae99819b 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
@@ -7,6 +7,7 @@ module ActiveRecord
         def initialize(klass, records, reflection, preload_options)
           super
           @join_table = Arel::Table.new(reflection.join_table).alias('t0')
+          @records_by_owner = nil
         end
 
         # Unlike the other associations, we want to get a raw array of rows so that we can
@@ -34,8 +35,10 @@ module ActiveRecord
         # actual records, ensuring that we don't create more than one instances of the same
         # record
         def associated_records_by_owner
+          return @records_by_owner if @records_by_owner
+
           records = {}
-          super.each_value do |rows|
+          @records_by_owner = super.each_value do |rows|
             rows.map! { |row| records[row[klass.primary_key]] ||= klass.instantiate(row) }
           end
         end