Improved Memoizable test coverage and added support for multiple arguments

3 files changed, 165 insertions, 61 deletions

diff --git a/activesupport/lib/active_support/core_ext/object/metaclass.rb b/activesupport/lib/active_support/core_ext/object/metaclass.rb
index 169a76dfb7..93fb0ad594 100644
--- a/activesupport/lib/active_support/core_ext/object/metaclass.rb
+++ b/activesupport/lib/active_support/core_ext/object/metaclass.rb
@@ -5,4 +5,9 @@ class Object
       self
     end
   end
+
+  # If class_eval is called on an object, add those methods to its metaclass
+  def class_eval(*args, &block)
+    metaclass.class_eval(*args, &block)
+  end
 end
diff --git a/activesupport/lib/active_support/memoizable.rb b/activesupport/lib/active_support/memoizable.rb
index f7cd73d39c..21636b8af4 100644
--- a/activesupport/lib/active_support/memoizable.rb
+++ b/activesupport/lib/active_support/memoizable.rb
@@ -1,32 +1,43 @@
 module ActiveSupport
-  module Memoizable #:nodoc:
+  module Memoizable
+    module Freezable
+      def self.included(base)
+        base.class_eval do
+          unless base.method_defined?(:freeze_without_memoizable)
+            alias_method_chain :freeze, :memoizable
+          end
+        end
+      end
+
+      def freeze_with_memoizable
+        methods.each do |method|
+          if m = method.to_s.match(/^_unmemoized_(.*)/)
+            send(m[1])
+          end
+        end
+        freeze_without_memoizable
+      end
+    end
+
     def memoize(*symbols)
       symbols.each do |symbol|
-        original_method = "unmemoized_#{symbol}"
-        memoized_ivar = "@#{symbol}"
+        original_method = "_unmemoized_#{symbol}"
+        memoized_ivar = "@_memoized_#{symbol}"
 
-        klass = respond_to?(:class_eval) ? self : self.metaclass
-        raise "Already memoized #{symbol}" if klass.instance_methods.map(&:to_s).include?(original_method)
+        class_eval <<-EOS, __FILE__, __LINE__
+          include Freezable
 
-        klass.class_eval <<-EOS, __FILE__, __LINE__
-          unless instance_methods.map(&:to_s).include?("freeze_without_memoizable")
-            alias_method :freeze_without_memoizable, :freeze
-            def freeze
-              methods.each do |method|
-                if m = method.to_s.match(/^unmemoized_(.*)/)
-                  send(m[1])
-                end
-              end
-              freeze_without_memoizable
-            end
-          end
+          raise "Already memoized #{symbol}" if method_defined?(:#{original_method})
+          alias #{original_method} #{symbol}
+
+          def #{symbol}(*args)
+            #{memoized_ivar} ||= {}
+            reload = args.pop if args.last == true || args.last == :reload
 
-          alias_method :#{original_method}, :#{symbol}
-          def #{symbol}(reload = false)
-            if !reload && defined? #{memoized_ivar}
-              #{memoized_ivar}
+            if !reload && #{memoized_ivar} && #{memoized_ivar}.has_key?(args)
+              #{memoized_ivar}[args]
             else
-              #{memoized_ivar} = #{original_method}.freeze
+              #{memoized_ivar}[args] = #{original_method}(*args).freeze
             end
           end
         EOS
diff --git a/activesupport/test/memoizable_test.rb b/activesupport/test/memoizable_test.rb
index 79769631ad..cd84dcda53 100644
--- a/activesupport/test/memoizable_test.rb
+++ b/activesupport/test/memoizable_test.rb
@@ -5,86 +5,174 @@ uses_mocha 'Memoizable' do
     class Person
       extend ActiveSupport::Memoizable
 
-      def name
-        fetch_name_from_floppy
+      attr_reader :name_calls, :age_calls
+      def initialize
+        @name_calls = 0
+        @age_calls = 0
       end
 
-      memoize :name
+      def name
+        @name_calls += 1
+        "Josh"
+      end
 
       def age
+        @age_calls += 1
         nil
       end
 
-      def counter
-        @counter ||= 0
-        @counter += 1
+      memoize :name, :age
+    end
+
+    class Company
+      attr_reader :name_calls
+      def initialize
+        @name_calls = 0
       end
 
-      memoize :age, :counter
+      def name
+        @name_calls += 1
+        "37signals"
+      end
+    end
+
+    module Rates
+      extend ActiveSupport::Memoizable
 
-      private
-        def fetch_name_from_floppy
-          "Josh"
+      attr_reader :sales_tax_calls
+      def sales_tax(price)
+        @sales_tax_calls ||= 0
+        @sales_tax_calls += 1
+        price * 0.1025
+      end
+      memoize :sales_tax
+    end
+
+    class Calculator
+      extend ActiveSupport::Memoizable
+      include Rates
+
+      attr_reader :fib_calls
+      def initialize
+        @fib_calls = 0
+      end
+
+      def fib(n)
+        @fib_calls += 1
+
+        if n == 0 || n == 1
+          n
+        else
+          fib(n - 1) + fib(n - 2)
         end
+      end
+      memoize :fib
+
+      def counter
+        @count ||= 0
+        @count += 1
+      end
+      memoize :counter
     end
 
     def setup
       @person = Person.new
+      @calculator = Calculator.new
     end
 
     def test_memoization
       assert_equal "Josh", @person.name
+      assert_equal 1, @person.name_calls
 
-      @person.expects(:fetch_name_from_floppy).never
-      2.times { assert_equal "Josh", @person.name }
+      3.times { assert_equal "Josh", @person.name }
+      assert_equal 1, @person.name_calls
+    end
+
+    def test_memoization_with_nil_value
+      assert_equal nil, @person.age
+      assert_equal 1, @person.age_calls
+
+      3.times { assert_equal nil, @person.age }
+      assert_equal 1, @person.age_calls
     end
 
     def test_reloadable
-      counter = @person.counter
-      assert_equal 1, @person.counter
-      assert_equal 2, @person.counter(:reload)
+      counter = @calculator.counter
+      assert_equal 1, @calculator.counter
+      assert_equal 2, @calculator.counter(:reload)
+      assert_equal 2, @calculator.counter
+      assert_equal 3, @calculator.counter(true)
+      assert_equal 3, @calculator.counter
     end
 
-    def test_memoized_methods_are_frozen
-      assert_equal true, @person.name.frozen?
+    def test_memoization_cache_is_different_for_each_instance
+      assert_equal 1, @calculator.counter
+      assert_equal 2, @calculator.counter(:reload)
+      assert_equal 1, Calculator.new.counter
+    end
 
+    def test_memoized_is_not_affected_by_freeze
       @person.freeze
       assert_equal "Josh", @person.name
-      assert_equal true, @person.name.frozen?
     end
 
-    def test_memoization_frozen_with_nil_value
-      @person.freeze
-      assert_equal nil, @person.age
+    def test_memoization_with_args
+      assert_equal 55, @calculator.fib(10)
+      assert_equal 11, @calculator.fib_calls
     end
 
-    def test_double_memoization
-      assert_raise(RuntimeError) { Person.memoize :name }
+    def test_reloadable_with_args
+      assert_equal 55, @calculator.fib(10)
+      assert_equal 11, @calculator.fib_calls
+      assert_equal 55, @calculator.fib(10, :reload)
+      assert_equal 12, @calculator.fib_calls
+      assert_equal 55, @calculator.fib(10, true)
+      assert_equal 13, @calculator.fib_calls
     end
 
-    class Company
-      def name
-        lookup_name
+    def test_object_memoization
+      [Company.new, Company.new, Company.new].each do |company|
+        company.extend ActiveSupport::Memoizable
+        company.memoize :name
+
+        assert_equal "37signals", company.name
+        assert_equal 1, company.name_calls
+        assert_equal "37signals", company.name
+        assert_equal 1, company.name_calls
       end
+    end
 
-      def lookup_name
-        "37signals"
-      end
+    def test_memoized_module_methods
+      assert_equal 1.025, @calculator.sales_tax(10)
+      assert_equal 1, @calculator.sales_tax_calls
+      assert_equal 1.025, @calculator.sales_tax(10)
+      assert_equal 1, @calculator.sales_tax_calls
+      assert_equal 2.5625, @calculator.sales_tax(25)
+      assert_equal 2, @calculator.sales_tax_calls
     end
 
-    def test_object_memoization
+    def test_object_memoized_module_methods
       company = Company.new
-      company.extend ActiveSupport::Memoizable
-      company.memoize :name
+      company.extend(Rates)
+
+      assert_equal 1.025, company.sales_tax(10)
+      assert_equal 1, company.sales_tax_calls
+      assert_equal 1.025, company.sales_tax(10)
+      assert_equal 1, company.sales_tax_calls
+      assert_equal 2.5625, company.sales_tax(25)
+      assert_equal 2, company.sales_tax_calls
+    end
 
-      assert_equal "37signals", company.name
-      # Mocha doesn't play well with frozen objects
-      company.metaclass.instance_eval { define_method(:lookup_name) { b00m } }
-      assert_equal "37signals", company.name
+    def test_double_memoization
+      assert_raise(RuntimeError) { Person.memoize :name }
+      person = Person.new
+      person.extend ActiveSupport::Memoizable
+      assert_raise(RuntimeError) { person.memoize :name }
 
-      assert_equal true, company.name.frozen?
-      company.freeze
-      assert_equal true, company.name.frozen?
+      company = Company.new
+      company.extend ActiveSupport::Memoizable
+      company.memoize :name
+      assert_raise(RuntimeError) { company.memoize :name }
     end
   end
 end