Faster2: Cache the method, ivar, and arity and the ancestry memoized methods

lib/memoist.rb

@@ -3,7 +3,7 @@
 module Memoist
 
   def self.memoized_ivar_for(method_name, identifier=nil)
-    "@#{memoized_prefix(identifier)}_#{escape_punctuation(method_name.to_s)}"
+    "@#{memoized_prefix(identifier)}_#{escape_punctuation(method_name)}"
   end
 
   def self.unmemoized_method_for(method_name, identifier=nil)
@@ -27,9 +27,9 @@ def self.unmemoized_prefix(identifier=nil)
   end
 
   def self.escape_punctuation(string)
-    return string unless string.end_with?('?'.freeze, '!'.freeze)
+    string = string.is_a?(String) ? string.dup : string.to_s
 
-    string = string.dup
+    return string unless string.end_with?('?'.freeze, '!'.freeze)
 
     # A String can't end in both ? and !
     if string.sub!(/\?\Z/, '_query'.freeze)
@@ -63,37 +63,60 @@ def unmemoize_all
       flush_cache
     end
 
+    def memoized_structs(names)
+      structs = self.class.all_memoized_structs
+      return structs if names.empty?
+
+      structs.select { |s| names.include?(s.memoized_method) }
+    end
+
     def prime_cache(*method_names)
-      method_names = self.class.memoized_methods if method_names.empty?
-      method_names.each do |method_name|
-        if method(Memoist.unmemoized_method_for(method_name)).arity == 0
-          __send__(method_name)
+      memoized_structs(method_names).each do |struct|
+        if struct.arity == 0
+          __send__(struct.memoized_method)
         else
-          ivar = Memoist.memoized_ivar_for(method_name)
-          instance_variable_set(ivar, {})
+          instance_variable_set(struct.ivar, {})
         end
       end
     end
 
     def flush_cache(*method_names)
-      method_names = self.class.memoized_methods if method_names.empty?
+      memoized_structs(method_names).each do |struct|
+        remove_instance_variable(struct.ivar) if instance_variable_defined?(struct.ivar)
+      end
+    end
+  end
+
+  MemoizedMethod = Struct.new(:memoized_method, :ivar, :arity)
+
+  def all_memoized_structs
+    @all_memoized_structs ||= begin
+      structs = memoized_methods.dup
 
-      method_names.each do |method_name|
-        ivar = Memoist.memoized_ivar_for(method_name)
-        remove_instance_variable(ivar) if instance_variable_defined?(ivar)
+      # Collect the memoized_methods of ancestors in ancestor order
+      # unless we already have it since self or parents could be overriding
+      # an ancestor method.
+      ancestors.grep(Memoist).each do |ancestor|
+        ancestor.memoized_methods.each do |m|
+          structs << m unless structs.any? {|am| am.memoized_method == m.memoized_method }
+        end
       end
+      structs
     end
   end
 
+  def clear_structs
+    @all_memoized_structs = nil
+  end
+
   def memoize(*method_names)
     if method_names.last.is_a?(Hash)
       identifier = method_names.pop[:identifier]
     end
 
     Memoist.memoist_eval(self) do
       def self.memoized_methods
-        require 'set'
-        @_memoized_methods ||= Set.new
+        @_memoized_methods ||= []
       end
     end
 
@@ -110,8 +133,9 @@ def self.memoized_methods
         end
         alias_method unmemoized_method, method_name
 
-        self.memoized_methods << method_name
-        if instance_method(method_name).arity == 0
+        mm = MemoizedMethod.new(method_name, memoized_ivar, instance_method(method_name).arity)
+        self.memoized_methods << mm
+        if mm.arity == 0
 
           # define a method like this;
 

test/memoist_test.rb

@@ -114,6 +114,13 @@ def name
     memoize :name, :identifier => :student
   end
 
+  class Teacher < Person
+    def seniority
+      "very_senior"
+    end
+    memoize :seniority
+  end
+
   class Company
     attr_reader :name_calls
     def initialize
@@ -261,11 +268,75 @@ def test_unmemoize_all
     assert_equal 2, @calculator.counter
   end
 
+  def test_all_memoized_structs
+    # Person             memoize :age, :is_developer?, :memoize_protected_test, :name, :name?, :sleep, :update, :update_attributes
+    # Student < Person   memoize :name, :identifier => :student
+    # Teacher < Person   memoize :seniority
+
+    expected = %w(age is_developer? memoize_protected_test name name? sleep update update_attributes)
+    structs = Person.all_memoized_structs
+    assert_equal expected, structs.collect(&:memoized_method).collect(&:to_s).sort
+    assert_equal "@_memoized_name", structs.detect {|s| s.memoized_method == :name }.ivar
+
+    # Same expected methods
+    structs = Student.all_memoized_structs
+    assert_equal expected, structs.collect(&:memoized_method).collect(&:to_s).sort
+    assert_equal "@_memoized_student_name", structs.detect {|s| s.memoized_method == :name }.ivar
+
+    expected = (expected << "seniority").sort
+    structs = Teacher.all_memoized_structs
+    assert_equal expected, structs.collect(&:memoized_method).collect(&:to_s).sort
+    assert_equal "@_memoized_name", structs.detect {|s| s.memoized_method == :name }.ivar
+  end
+
+  def test_unmemoize_all_subclasses
+    # Person             memoize :age, :is_developer?, :memoize_protected_test, :name, :name?, :sleep, :update, :update_attributes
+    # Student < Person   memoize :name, :identifier => :student
+    # Teacher < Person   memoize :seniority
+
+    teacher = Teacher.new
+    assert_equal "Josh", teacher.name
+    assert_equal "Josh", teacher.instance_variable_get(:@_memoized_name)
+    assert_equal "very_senior", teacher.seniority
+    assert_equal "very_senior", teacher.instance_variable_get(:@_memoized_seniority)
+
+    teacher.unmemoize_all
+    assert_nil teacher.instance_variable_get(:@_memoized_name)
+    assert_nil teacher.instance_variable_get(:@_memoized_seniority)
+
+    student = Student.new
+    assert_equal "Student Josh", student.name
+    assert_equal "Student Josh", student.instance_variable_get(:@_memoized_student_name)
+    assert_nil student.instance_variable_get(:@_memoized_seniority)
+
+    student.unmemoize_all
+    assert_nil student.instance_variable_get(:@_memoized_student_name)
+  end
+
   def test_memoize_all
     @calculator.memoize_all
     assert @calculator.instance_variable_defined?(:@_memoized_counter)
   end
 
+  def test_memoize_all_subclasses
+    # Person             memoize :age, :is_developer?, :memoize_protected_test, :name, :name?, :sleep, :update, :update_attributes
+    # Student < Person   memoize :name, :identifier => :student
+    # Teacher < Person   memoize :seniority
+
+    teacher = Teacher.new
+    teacher.memoize_all
+
+    assert_equal "very_senior", teacher.instance_variable_get(:@_memoized_seniority)
+    assert_equal "Josh", teacher.instance_variable_get(:@_memoized_name)
+
+    student = Student.new
+    student.memoize_all
+
+    assert_equal "Student Josh", student.instance_variable_get(:@_memoized_student_name)
+    assert_equal "Student Josh", student.name
+    assert_nil student.instance_variable_get(:@_memoized_seniority)
+  end
+
   def test_memoization_cache_is_different_for_each_instance
     assert_equal 1, @calculator.counter
     assert_equal 2, @calculator.counter(:reload)
@@ -331,7 +402,29 @@ def test_object_memoized_module_methods
   end
 
   def test_double_memoization_with_identifier
+    # Person             memoize :age, :is_developer?, :memoize_protected_test, :name, :name?, :sleep, :update, :update_attributes
+    # Student < Person   memoize :name, :identifier => :student
+    # Teacher < Person   memoize :seniority
+
     Person.memoize :name, :identifier => :again
+    p = Person.new
+    assert_equal "Josh", p.name
+    assert p.instance_variable_get(:@_memoized_again_name)
+
+    # HACK: tl;dr: Don't memoize classes in test that are used elsewhere.
+    # Calling Person.memoize :name, :identifier => :again pollutes Person
+    # and descendents since we cache the memoized method structures.
+    # This populates those structs, verifies Person is polluted, resets the
+    # structs, cleans up cached memoized_methods
+    Student.all_memoized_structs
+    Person.all_memoized_structs
+    Teacher.all_memoized_structs
+    assert Person.memoized_methods.any? { |m| m.ivar == "@_memoized_again_name" }
+
+    [Student, Teacher, Person].each { |obj| obj.clear_structs }
+    assert Person.memoized_methods.reject!      { |m| m.ivar == "@_memoized_again_name" }
+    assert_nil Student.memoized_methods.reject! { |m| m.ivar == "@_memoized_again_name" }
+    assert_nil Teacher.memoized_methods.reject! { |m| m.ivar == "@_memoized_again_name" }
   end
 
   def test_memoization_with_a_subclass