lib/memoist.rb in memoist-0.13.0 vs lib/memoist.rb in memoist-0.14.0

@@ -1,11 +1,11 @@
 require 'memoist/core_ext/singleton_class'
 
 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)
     "#{unmemoized_prefix(identifier)}_#{method_name}".to_sym
   end
@@ -25,14 +25,14 @@
       "_unmemoized".freeze
     end
   end
 
   def self.escape_punctuation(string)
+    string = string.is_a?(String) ? string.dup : string.to_s
+
     return string unless string.end_with?('?'.freeze, '!'.freeze)
 
-    string = string.dup
-
     # A String can't end in both ? and !
     if string.sub!(/\?\Z/, '_query'.freeze)
     else
       string.sub!(/!\Z/, '_bang'.freeze)
     end
@@ -61,41 +61,64 @@
 
     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.dup
+
+      # 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
 
     method_names.each do |method_name|
       unmemoized_method = Memoist.unmemoized_method_for(method_name, identifier)
@@ -108,11 +131,12 @@
           warn "Already memoized #{method_name}"
           return
         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;
 
           # def mime_type(reload=true)
           #   skip_cache = reload || !instance_variable_defined?("@_memoized_mime_type")