IdentityMap - Adding Weakling and IM Base as concern

2 files changed, 66 insertions, 0 deletions

diff --git a/activerecord/lib/active_record.rb b/activerecord/lib/active_record.rb
index c80bce2849..ff6cf44a73 100644
--- a/activerecord/lib/active_record.rb
+++ b/activerecord/lib/active_record.rb
@@ -32,6 +32,7 @@ require 'active_support'
 require 'active_support/i18n'
 require 'active_model'
 require 'arel'
+require 'weakling'
 
 require 'active_record/version'
 
@@ -79,6 +80,7 @@ module ActiveRecord
     autoload :Timestamp
     autoload :Transactions
     autoload :Validations
+    autoload :IdentityMap
   end
 
   module AttributeMethods
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
new file mode 100644
index 0000000000..79eb4dd83a
--- /dev/null
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -0,0 +1,64 @@
+module ActiveRecord
+  module IdentityMap
+    extend ActiveSupport::Concern
+
+    class << self
+      attr_accessor :repositories
+      attr_accessor :current_repository_name
+      attr_accessor :enabled
+
+      def current
+        repositories[current_repository_name] ||= Weakling::WeakHash.new
+      end
+
+      def with_repository(name = :default, &block)
+        old_repository = self.current_repository_name
+        self.current_repository_name = name
+
+        block.call(current)
+      ensure
+        self.current_repository_name = old_repository
+      end
+
+      def without(&block)
+        old, self.enabled = self.enabled, false
+
+        block.call
+      ensure
+        self.enabled = old
+      end
+
+      def get(class_name, primary_key)
+        current[[class_name, primary_key]]
+      end
+
+      def add(record)
+        current[[record.class.name, record.id]] = record
+      end
+
+      def remove(record)
+        current.delete([record.class.name, record.id])
+      end
+
+      def clear
+        current.clear
+      end
+
+      alias enabled? enabled
+    end
+
+    self.repositories ||= Hash.new
+    self.current_repository_name ||= :default
+    self.enabled = true
+
+    module InstanceMethods
+
+    end
+
+    module ClassMethods
+      def identity_map
+        ActiveRecord::IdentityMap
+      end
+    end
+  end
+end