1 files changed, 16 insertions, 1 deletions
diff --git a/activerecord/lib/active_record/associations/association_proxy.rb b/activerecord/lib/active_record/associations/association_proxy.rb
index fc03a4b619..d16cda5585 100644
--- a/activerecord/lib/active_record/associations/association_proxy.rb
+++ b/activerecord/lib/active_record/associations/association_proxy.rb
@@ -111,6 +111,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
 
@@ -228,7 +229,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
-          @target = find_target if find_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
           loaded!
           target
         rescue ActiveRecord::RecordNotFound
@@ -328,6 +339,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