1 files changed, 16 insertions, 0 deletions
diff --git a/activerecord/lib/active_record/associations/association.rb b/activerecord/lib/active_record/associations/association.rb
index 2264631584..ae745ea7c2 100644
--- a/activerecord/lib/active_record/associations/association.rb
+++ b/activerecord/lib/active_record/associations/association.rb
@@ -42,6 +42,7 @@ module ActiveRecord
       # Resets the \loaded flag to +false+ and sets the \target to +nil+.
       def reset
         @loaded = false
+        IdentityMap.remove(@target) if defined?(@target) && @target && IdentityMap.enabled?
         @target = nil
       end
 
@@ -141,6 +142,17 @@ module ActiveRecord
       # ActiveRecord::RecordNotFound is rescued within the method, and it is
       # not reraised. The proxy is \reset and +nil+ is the return value.
       def load_target
+        if find_target?
+          begin
+          if IdentityMap.enabled? && association_class && association_class.respond_to?(:base_class)
+            @target = IdentityMap.get(association_class, @owner[@reflection.foreign_key])
+          end
+          rescue NameError
+            nil
+          ensure
+            @target ||= find_target
+          end
+        end
         @target = find_target if find_target?
         loaded!
         target
@@ -241,6 +253,10 @@ module ActiveRecord
         # This is only relevant to certain associations, which is why it returns nil by default.
         def stale_state
         end
+
+        def association_class
+          @reflection.klass
+        end
     end
   end
 end