Fix preloader to never reset associations in case they are already loaded

This patch fixes the issue when association is preloaded with a custom
preload scope which disposes the already preloaded target of the
association by reseting it.

When custom preload scope is used, the preloading is now performed into
a separated Hash - #records_by_owner instead of the association.
It removes the necessaty the reset the association after the preloading
is complete so that reset of the preloaded association never happens.

Preloading is still happening to the association when the preload scope
is empty.

4 files changed, 99 insertions, 64 deletions

diff --git a/activerecord/lib/active_record/associations/preloader.rb b/activerecord/lib/active_record/associations/preloader.rb
index 8997579527..c7cd87f9d4 100644
--- a/activerecord/lib/active_record/associations/preloader.rb
+++ b/activerecord/lib/active_record/associations/preloader.rb
@@ -143,9 +143,7 @@ module ActiveRecord
 
         def preloaders_for_reflection(reflection, records, scope)
           records.group_by { |record| record.association(reflection.name).klass }.map do |rhs_klass, rs|
-            loader = preloader_for(reflection, rs).new(rhs_klass, rs, reflection, scope)
-            loader.run self
-            loader
+            preloader_for(reflection, rs).new(rhs_klass, rs, reflection, scope).run
           end
         end
 
@@ -166,10 +164,18 @@ module ActiveRecord
             @reflection = reflection
           end
 
-          def run(preloader); end
+          def run
+            self
+          end
 
           def preloaded_records
-            owners.flat_map { |owner| owner.association(reflection.name).target }
+            @preloaded_records ||= records_by_owner.flat_map(&:last)
+          end
+
+          def records_by_owner
+            @records_by_owner ||= owners.each_with_object({}) do |owner, result|
+              result[owner] = Array(owner.association(reflection.name).target)
+            end
           end
 
           private
diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index 7048ff43b8..46532f651e 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -4,26 +4,44 @@ module ActiveRecord
   module Associations
     class Preloader
       class Association #:nodoc:
-        attr_reader :preloaded_records
-
         def initialize(klass, owners, reflection, preload_scope)
           @klass         = klass
           @owners        = owners
           @reflection    = reflection
           @preload_scope = preload_scope
           @model         = owners.first && owners.first.class
-          @preloaded_records = []
         end
 
-        def run(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)
+        def run
+          if !preload_scope || preload_scope.empty_scope?
+            owners.each do |owner|
+              associate_records_to_owner(owner, records_by_owner[owner] || [])
+            end
+          else
+            # Custom preload scope is used and
+            # the association can not be marked as loaded
+            # Loading into a Hash instead
+            records_by_owner
           end
+          self
+        end
 
-          owners.each do |owner|
-            associate_records_to_owner(owner, records[convert_key(owner[owner_key_name])] || [])
+        def records_by_owner
+          @records_by_owner ||= preloaded_records.each_with_object({}) do |record, result|
+            owners_by_key[convert_key(record[association_key_name])].each do |owner|
+              (result[owner] ||= []) << record
+            end
+          end
+        end
+
+        def preloaded_records
+          return @preloaded_records if defined?(@preloaded_records)
+          return [] if owner_keys.empty?
+          # 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
+          slices = owner_keys.each_slice(klass.connection.in_clause_length || owner_keys.size)
+          @preloaded_records = slices.flat_map do |slice|
+            records_for(slice)
           end
         end
 
@@ -54,13 +72,10 @@ module ActiveRecord
           end
 
           def owners_by_key
-            unless defined?(@owners_by_key)
-              @owners_by_key = owners.each_with_object({}) do |owner, h|
-                key = convert_key(owner[owner_key_name])
-                h[key] = owner if key
-              end
+            @owners_by_key ||= owners.each_with_object({}) do |owner, result|
+              key = convert_key(owner[owner_key_name])
+              (result[key] ||= []) << owner if key
             end
-            @owners_by_key
           end
 
           def key_conversion_required?
@@ -87,23 +102,16 @@ module ActiveRecord
             @model.type_for_attribute(owner_key_name).type
           end
 
-          def load_records(&block)
-            return {} if owner_keys.empty?
-            # 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
-            slices = owner_keys.each_slice(klass.connection.in_clause_length || owner_keys.size)
-            @preloaded_records = slices.flat_map do |slice|
-              records_for(slice, &block)
-            end
-            @preloaded_records.group_by do |record|
-              convert_key(record[association_key_name])
+          def records_for(ids)
+            scope.where(association_key_name => ids).load do |record|
+              # Processing only the first owner
+              # because the record is modified but not an owner
+              owner = owners_by_key[convert_key(record[association_key_name])].first
+              association = owner.association(reflection.name)
+              association.set_inverse_instance(record)
             end
           end
 
-          def records_for(ids, &block)
-            scope.where(association_key_name => ids).load(&block)
-          end
-
           def scope
             @scope ||= build_scope
           end
diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb
index 32653956b2..54b019e366 100644
--- a/activerecord/lib/active_record/associations/preloader/through_association.rb
+++ b/activerecord/lib/active_record/associations/preloader/through_association.rb
@@ -4,45 +4,57 @@ module ActiveRecord
   module Associations
     class Preloader
       class ThroughAssociation < Association # :nodoc:
-        def run(preloader)
-          already_loaded     = owners.first.association(through_reflection.name).loaded?
-          through_scope      = through_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, scope)
-          @preloaded_records = preloaders.flat_map(&:preloaded_records)
-
-          owners.each do |owner|
-            through_records = Array(owner.association(through_reflection.name).target)
-
-            if already_loaded
+        PRELOADER = ActiveRecord::Associations::Preloader.new
+
+        def initialize(*)
+          super
+          @already_loaded = owners.first.association(through_reflection.name).loaded?
+        end
+
+        def preloaded_records
+          @preloaded_records ||= source_preloaders.flat_map(&:preloaded_records)
+        end
+
+        def records_by_owner
+          return @records_by_owner if defined?(@records_by_owner)
+          source_records_by_owner = source_preloaders.map(&:records_by_owner).reduce(:merge)
+          through_records_by_owner = through_preloaders.map(&:records_by_owner).reduce(:merge)
+
+          @records_by_owner = owners.each_with_object({}) do |owner, result|
+            through_records = through_records_by_owner[owner] || []
+
+            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|
-              record.association(source_reflection.name).target
+            records = through_records.flat_map do |record|
+              source_records_by_owner[record]
             end
 
-            result.compact!
-            result.sort_by! { |rhs| preload_index[rhs] } if scope.order_values.any?
-            result.uniq! if scope.distinct_value
-            associate_records_to_owner(owner, result)
-          end
-
-          unless scope.empty_scope?
-            middle_records.each do |owner|
-              owner.association(source_reflection.name).reset if owner
-            end
+            records.compact!
+            records.sort_by! { |rhs| preload_index[rhs] } if scope.order_values.any?
+            records.uniq! if scope.distinct_value
+            result[owner] = records
           end
         end
 
         private
+          def source_preloaders
+            @source_preloaders ||= PRELOADER.preload(middle_records, source_reflection.name, scope)
+          end
+
+          def middle_records
+            through_preloaders.flat_map(&:preloaded_records)
+          end
+
+          def through_preloaders
+            @through_preloaders ||= PRELOADER.preload(owners, through_reflection.name, through_scope)
+          end
+
           def through_reflection
             reflection.through_reflection
           end
@@ -52,8 +64,8 @@ module ActiveRecord
           end
 
           def preload_index
-            @preload_index ||= @preloaded_records.each_with_object({}).with_index do |(id, result), index|
-              result[id] = index
+            @preload_index ||= preloaded_records.each_with_object({}).with_index do |(record, result), index|
+              result[record] = index
             end
           end
 
@@ -92,7 +104,7 @@ module ActiveRecord
               end
             end
 
-            scope unless scope.empty_scope?
+            scope
           end
       end
     end
diff --git a/activerecord/test/cases/associations/nested_through_associations_test.rb b/activerecord/test/cases/associations/nested_through_associations_test.rb
index 0b83fd8421..35da74102d 100644
--- a/activerecord/test/cases/associations/nested_through_associations_test.rb
+++ b/activerecord/test/cases/associations/nested_through_associations_test.rb
@@ -548,6 +548,15 @@ class NestedThroughAssociationsTest < ActiveRecord::TestCase
     end
   end
 
+  def test_through_association_preload_doesnt_reset_source_association_if_already_preloaded
+    blue = tags(:blue)
+    authors = Author.preload(posts: :first_blue_tags_2, misc_post_first_blue_tags_2: {}).to_a.sort_by(&:id)
+
+    assert_no_queries do
+      assert_equal [blue], authors[2].posts.first.first_blue_tags_2
+    end
+  end
+
   def test_nested_has_many_through_with_conditions_on_source_associations_preload_via_joins
     # Pointless condition to force single-query loading
     assert_includes_and_joins_equal(