Adding base method symbolized_sti_name to activerecord base to be used on identity map. Identity map now considers the inheritance when creating the caching keys

3 files changed, 15 insertions, 9 deletions

diff --git a/activerecord/lib/active_record/base.rb b/activerecord/lib/active_record/base.rb
index 58a056bce9..32edae4ded 100644
--- a/activerecord/lib/active_record/base.rb
+++ b/activerecord/lib/active_record/base.rb
@@ -830,6 +830,10 @@ module ActiveRecord #:nodoc:
         @symbolized_base_class ||= base_class.to_s.to_sym
       end
 
+      def symbolized_sti_name
+        @symbolized_sti_name ||= sti_name ? sti_name.to_sym : symbolized_base_class
+      end
+
       # Returns the base AR subclass that this class descends from. If A
       # extends AR::Base, A.base_class will return A. If B descends from A
       # through some arbitrarily deep hierarchy, B.base_class will return A.
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 9eb47ad99f..f88ead9ca0 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -49,7 +49,7 @@ module ActiveRecord
       end
 
       def get(klass, primary_key)
-        record = repository[klass.symbolized_base_class][primary_key]
+        record = repository[klass.symbolized_sti_name][primary_key]
 
         if record.is_a?(klass)
           ActiveSupport::Notifications.instrument("identity.active_record",
@@ -64,15 +64,15 @@ module ActiveRecord
       end
 
       def add(record)
-        repository[record.class.symbolized_base_class][record.id] = record
+        repository[record.class.symbolized_sti_name][record.id] = record
       end
 
       def remove(record)
-        repository[record.class.symbolized_base_class].delete(record.id)
+        repository[record.class.symbolized_sti_name].delete(record.id)
       end
 
-      def remove_by_id(symbolized_base_class, id)
-        repository[symbolized_base_class].delete(id)
+      def remove_by_id(symbolized_sti_name, id)
+        repository[symbolized_sti_name].delete(id)
       end
 
       def clear
diff --git a/activerecord/test/cases/identity_map_test.rb b/activerecord/test/cases/identity_map_test.rb
index c183690187..a0e16400d2 100644
--- a/activerecord/test/cases/identity_map_test.rb
+++ b/activerecord/test/cases/identity_map_test.rb
@@ -143,8 +143,9 @@ class IdentityMapTest < ActiveRecord::TestCase
   
   def test_inherited_with_type_attribute_without_identity_map
     ActiveRecord::IdentityMap.without do
-      c1 = comments(:sub_special_comment)
-      c2 = Comment.find(c1.id)
+      c = comments(:sub_special_comment)
+      c1 = SubSpecialComment.find(c.id)
+      c2 = Comment.find(c.id)
       assert_same(c1.class, c2.class)
     end
   end
@@ -156,8 +157,9 @@ class IdentityMapTest < ActiveRecord::TestCase
   end
 
   def test_inherited_with_type_attribute
-    c1 = comments(:sub_special_comment)
-    c2 = Comment.find(c1.id)
+    c = comments(:sub_special_comment)
+    c1 = SubSpecialComment.find(c.id)
+    c2 = Comment.find(c.id)
     assert_same(c1, c2)
   end