active-support-memoization-prime-on-freeze.patch

active-support-memoization-prime-on-freeze.patch

diff --git a/activesupport/lib/active_support/memoizable.rb b/activesupport/lib/active_support/memoizable.rb
index 0a7bcd5..71af70c 100644
--- a/activesupport/lib/active_support/memoizable.rb
+++ b/activesupport/lib/active_support/memoizable.rb
@@ -17,7 +17,7 @@ module ActiveSupport
       end 
 
       def freeze_with_memoizable
-        memoize_all unless frozen?
+        prime_cache(*self.class.instance_variable_get(:"@prime_on_freeze").to_a || []) unless frozen?
         freeze_without_memoizable
       end 
 
@@ -57,8 +57,15 @@ module ActiveSupport
     end 
 
     def memoize(*symbols)
+      on_freeze = (symbols.pop if symbols.last.is_a? Hash) || {:on_freeze => :prime}
       symbols.each do |symbol|
         original_method = :"_unmemoized_#{symbol}"
+        @prime_on_freeze ||= Set.new
+        case on_freeze[:on_freeze]
+        when :prime then @prime_on_freeze << symbol
+        when :flush then @prime_on_freeze.delete symbol
+        else raise "Illegal option: to memoize: on_freeze should => :prime, or :flush"
+        end
         memoized_ivar = ActiveSupport::Memoizable.memoized_ivar_for(symbol)
 
         class_eval <<-EOS, __FILE__, __LINE__ + 1 
diff --git a/activesupport/test/memoizable_test.rb b/activesupport/test/memoizable_test.rb
index bceac13..cc76fe5 100644
--- a/activesupport/test/memoizable_test.rb
+++ b/activesupport/test/memoizable_test.rb
@@ -36,6 +36,15 @@ class MemoizableTest < ActiveSupport::TestCase
 
     memoize :name, :age
 
+    def flush_on_freeze
+      'flush'
+    end
+    memoize :flush_on_freeze, :on_freeze => :flush
+    def prime_on_freeze
+      'prime'
+    end
+    memoize :prime_on_freeze, :on_freeze => :prime
+
     protected
 
     def memoize_protected_test

@@ -185,6 +194,14 @@ class MemoizableTest < ActiveSupport::TestCase
     assert_equal "Joshua", @person.update("Joshua")
   end

+  def test_does_not_memoize_marked_methods_on_freeze
+    @person.freeze
+    assert_raise(RuntimeError) { @person.flush_on_freeze }
+    @person.prime_on_freeze
+    assert_equal @person.instance_variable_defined?(:@_memoized_flush_on_freeze), false
+    assert @person.instance_variable_defined?(:@_memoized_prime_on_freeze)
+  end
+
   def test_memoization_with_args
     assert_equal 55, @calculator.fib(10)
     assert_equal 11, @calculator.fib_calls