IdentityMap - Adding Weakling and IM Base as concern

.gitignore

@@ -20,3 +20,4 @@ railties/test/initializer/root/log
 railties/doc
 railties/guides/output
 railties/tmp
+nbproject

Gemfile

@@ -25,6 +25,8 @@ gem "memcache-client", ">= 1.8.5"
 # AM
 gem "text-format", "~> 1.0.0"
 
+gem "weakling", :git => "git://github.com/swistak/weakling.git"
+
 platforms :mri_18 do
   gem "system_timer"
   gem "ruby-debug", ">= 0.10.3"

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

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