IdentityMap - adding and removing of records on create/update

2 files changed, 43 insertions, 5 deletions

diff --git a/activerecord/lib/active_record/base.rb b/activerecord/lib/active_record/base.rb
index fe660c1587..d7e58ba978 100644
--- a/activerecord/lib/active_record/base.rb
+++ b/activerecord/lib/active_record/base.rb
@@ -878,13 +878,41 @@ module ActiveRecord #:nodoc:
           finder_needs_type_condition? ? @relation.where(type_condition) : @relation
         end
 
+        def instantiate_without_im(object, record)
+          object.instance_variable_set(:@attributes, record)
+          object.instance_variable_set(:@attributes_cache, {})
+          object.instance_variable_set(:@new_record, false)
+          object.instance_variable_set(:@readonly, false)
+          object.instance_variable_set(:@destroyed, false)
+          object.instance_variable_set(:@marked_for_destruction, false)
+          object.instance_variable_set(:@previously_changed, {})
+          object.instance_variable_set(:@changed_attributes, {})
+
+          object
+        end
+
         # Finder methods must instantiate through this method to work with the
         # single-table inheritance model that makes it possible to create
         # objects of different types from the same table.
         def instantiate(record)
-          model = find_sti_class(record[inheritance_column]).allocate
-          model.init_with('attributes' => record)
-          model
+          sti_class = find_sti_class(record[inheritance_column])
+          record_id = sti_class.primary_key && record[sti_class.primary_key]
+          if ActiveRecord::IdentityMap.enabled? && record_id
+            if object = identity_map.get(sti_class.name, record_id)
+              object.instance_variable_get("@attributes").update(record)
+              object.instance_variable_get("@attributes_cache").replace({})
+            else
+              object = instantiate_without_im(sti_class.allocate, record)
+              identity_map.add(object)
+            end
+          else
+            object = instantiate_without_im(sti_class.allocate, record)
+          end
+
+          object.send(:_run_find_callbacks)
+          object.send(:_run_initialize_callbacks)
+
+          object
         end
 
         def find_sti_class(type_name)
diff --git a/activerecord/lib/active_record/persistence.rb b/activerecord/lib/active_record/persistence.rb
index 594a2214bb..a7181fd79b 100644
--- a/activerecord/lib/active_record/persistence.rb
+++ b/activerecord/lib/active_record/persistence.rb
@@ -64,7 +64,10 @@ module ActiveRecord
     # callbacks, Observer methods, or any <tt>:dependent</tt> association
     # options, use <tt>#destroy</tt>.
     def delete
-      self.class.delete(id) if persisted?
+      if persisted?
+        self.class.delete(id)
+        IdentityMap.remove(self)
+      end
       @destroyed = true
       freeze
     end
@@ -73,6 +76,7 @@ module ActiveRecord
     # that no changes should be made (since they can't be persisted).
     def destroy
       if persisted?
+        IdentityMap.remove(self)
         self.class.unscoped.where(self.class.arel_table[self.class.primary_key].eq(id)).delete_all
       end
 
@@ -196,7 +200,12 @@ module ActiveRecord
     def reload(options = nil)
       clear_aggregation_cache
       clear_association_cache
-      @attributes.update(self.class.unscoped { self.class.find(self.id, options) }.instance_variable_get('@attributes'))
+
+      IdentityMap.without do
+        fresh_object = self.class.unscoped { self.class.find(self.id, options) }
+        @attributes.update(fresh_object.instance_variable_get('@attributes'))
+      end
+
       @attributes_cache = {}
       self
     end
@@ -270,6 +279,7 @@ module ActiveRecord
 
       self.id ||= new_id
 
+      IdentityMap.add(self)
       @persisted = true
       id
     end