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

We can limit calls to memoized_ivar_for/unmemoized_method_for and
related methods for things we've already calculated initially on the
memoize :some_method call.

Fixes prime_cache not priming inherited memoized methods, broken by #36
Fixes (un)memoize_all not clearing/priming subclass identifiers

flush_cache now:             121143.3 i/s
flush_cache PR 34:            16079.4 i/s - 7.53x slower
flush_cache master:            3718.0 i/s - 32.58x slower

flush_cache with args now:    43484.0 i/s
flush_cache with args PR 34:  17488.5 i/s - 2.49x slower
flush_cache with args master: 17279.3 i/s - 2.52x slower

prime_cache now:              55946.9 i/s
prime_cache PR 34:            12665.0 i/s - 4.42x slower
prime_cache master:            6057.2 i/s - 9.24x slower

prime_cache with args now:    30015.3 i/s
prime_cache with args PR 34:  14012.1 i/s - 2.14x slower
prime_cache with args master: 13914.9 i/s - 2.16x slower

For a class with 40 memoized methods, total allocations is also
greatly reduced where each method below is called 1,000 times:

method | PR 34 | master | now
------ | ----- | ----- | ---
prime_cache with args | 363421 | 164000 | 45004
prime_cache no args | 630535 | 168000 | 41000
flush_cache with args | 643361 | 164000 | 5000
flush_cache no args | 907607 | 168000 | 1000

SCRIPT:
```ruby
def log_all_allocations(key = :line)
  require 'allocation_tracer'

  trace_keys = %i{path type class} if [:path, :file].include?(key)
  trace_keys = %i{line path type class} if key == :line

  ObjectSpace::AllocationTracer.setup(trace_keys)
  return_from_yield = nil
  result = ObjectSpace::AllocationTracer.trace do
    return_from_yield = yield
    nil
  end

  puts "Total: #{result.values.inject(0) { |count, v| count += v[0]}}"
  return_from_yield
end

require './lib/memoist'
require 'set'

$methods = Set.new
class Person
  extend Memoist
  1.upto(20) do |n|
    method = define_method("test_#{n}".to_sym) {}
    memoize(method)
    $methods.add(method)
  end

  1.upto(20) do |n|
    method = define_method("test_#{n}?".to_sym) {}
    memoize(method)
    $methods.add(method)
  end
end

puts "prime_cache with args"
p = Person.new
log_all_allocations(:line) do
  1_000.times do
    p.prime_cache(*$methods)
  end
end

puts "prime_cache no args"
p = Person.new
log_all_allocations(:line) do
  1_000.times do
    p.prime_cache
  end
end

puts "flush_cache with args"
p = Person.new
log_all_allocations(:line) do
  1_000.times do
    p.flush_cache(*$methods)
  end
end

puts "flush_cache no args"
p = Person.new
log_all_allocations(:line) do
  1_000.times do
    p.flush_cache
  end
end
```

lib/memoist.rb

@@ -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
 
-      method_names.each do |method_name|
-        ivar = Memoist.memoized_ivar_for(method_name)
-        remove_instance_variable(ivar) if instance_variable_defined?(ivar)
+  MemoizedMethod = Struct.new(:memoized_method, :ivar, :arity)
+
+  def all_memoized_structs
+    @all_memoized_structs ||= begin
+      structs = memoized_methods
+
+      # 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,28 @@ 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
+    Teacher.all_memoized_structs
+    assert Person.memoized_methods.any? { |m| m.ivar == "@_memoized_again_name" }
+
+    ObjectSpace.each_object(Memoist) { |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