memoize_ and unmemoize_all

2 files changed, 48 insertions, 18 deletions

diff --git a/activesupport/lib/active_support/memoizable.rb b/activesupport/lib/active_support/memoizable.rb
index b5adc2330c..5aa73bce52 100644
--- a/activesupport/lib/active_support/memoizable.rb
+++ b/activesupport/lib/active_support/memoizable.rb
@@ -10,25 +10,37 @@ module ActiveSupport
       end
 
       def freeze_with_memoizable
-        unless frozen?
-          methods.each do |method|
-            if method.to_s =~ /^_unmemoized_(.*)/
-              begin
-                __send__($1).freeze
-              rescue ArgumentError
-              end
+        memoize_all unless frozen?
+        freeze_without_memoizable
+      end
+
+      def memoize_all
+        methods.each do |m|
+          if m.to_s =~ /^_unmemoized_(.*)/
+            if method(m).arity == 0
+              __send__($1)
+            else
+              ivar = :"@_memoized_#{$1}"
+              instance_variable_set(ivar, {})
             end
           end
         end
+      end
 
-        freeze_without_memoizable
+      def unmemoize_all
+        methods.each do |m|
+          if m.to_s =~ /^_unmemoized_(.*)/
+            ivar = :"@_memoized_#{$1}"
+            instance_variable_get(ivar).clear if instance_variable_defined?(ivar)
+          end
+        end
       end
     end
 
     def memoize(*symbols)
       symbols.each do |symbol|
-        original_method = "_unmemoized_#{symbol}"
-        memoized_ivar = "@_memoized_#{symbol.to_s.sub(/\?\Z/, '_query').sub(/!\Z/, '_bang')}"
+        original_method = :"_unmemoized_#{symbol}"
+        memoized_ivar = :"@_memoized_#{symbol.to_s.sub(/\?\Z/, '_query').sub(/!\Z/, '_bang')}"
 
         class_eval <<-EOS, __FILE__, __LINE__
           include Freezable
@@ -38,21 +50,24 @@ module ActiveSupport
 
           if instance_method(:#{symbol}).arity == 0
             def #{symbol}(reload = false)
-              if !reload && defined? #{memoized_ivar}
-                #{memoized_ivar}
-              else
-                #{memoized_ivar} = #{original_method}
+              if reload || !defined?(#{memoized_ivar}) || #{memoized_ivar}.empty?
+                #{memoized_ivar} = [#{original_method}]
               end
+              #{memoized_ivar}[0]
             end
           else
             def #{symbol}(*args)
-              #{memoized_ivar} ||= {}
+              #{memoized_ivar} ||= {} unless frozen?
               reload = args.pop if args.last == true || args.last == :reload
 
-              if !reload && #{memoized_ivar} && #{memoized_ivar}.has_key?(args)
-                #{memoized_ivar}[args]
+              if #{memoized_ivar}
+                if !reload && #{memoized_ivar}.has_key?(args)
+                  #{memoized_ivar}[args]
+                elsif #{memoized_ivar}
+                  #{memoized_ivar}[args] = #{original_method}(*args)
+                end
               else
-                #{memoized_ivar}[args] = #{original_method}(*args)
+                #{original_method}(*args)
               end
             end
           end
diff --git a/activesupport/test/memoizable_test.rb b/activesupport/test/memoizable_test.rb
index 01631dc5e1..99d9240ea1 100644
--- a/activesupport/test/memoizable_test.rb
+++ b/activesupport/test/memoizable_test.rb
@@ -119,6 +119,21 @@ uses_mocha 'Memoizable' do
       assert_equal 3, @calculator.counter
     end
 
+    def test_unmemoize_all
+      assert_equal 1, @calculator.counter
+
+      assert @calculator.instance_variable_get(:@_memoized_counter).any?
+      @calculator.unmemoize_all
+      assert @calculator.instance_variable_get(:@_memoized_counter).empty?
+
+      assert_equal 2, @calculator.counter
+    end
+
+    def test_memoize_all
+      @calculator.memoize_all
+      assert @calculator.instance_variable_defined?(:@_memoized_counter)
+    end
+
     def test_memoization_cache_is_different_for_each_instance
       assert_equal 1, @calculator.counter
       assert_equal 2, @calculator.counter(:reload)