From 027305139318844772c3e46f97b5e924981a03f6 Mon Sep 17 00:00:00 2001
From: Joe Burns <joseph.burns@gmail.com>
Date: Wed, 31 May 2017 12:55:37 -0400
Subject: [PATCH] Add support for class-level cache flushing.

---
 lib/memoist.rb       |  3 ++-
 test/memoist_test.rb | 58 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/lib/memoist.rb b/lib/memoist.rb
index 92e1a7b..1f6bb21 100644
--- a/lib/memoist.rb
+++ b/lib/memoist.rb
@@ -66,7 +66,8 @@ def unmemoize_all
     end
 
     def memoized_structs(names)
-      structs = self.class.all_memoized_structs
+      ref_obj = self.class.respond_to?(:class_eval) ? self.singleton_class : self
+      structs = ref_obj.all_memoized_structs
       return structs if names.empty?
 
       structs.select { |s| names.include?(s.memoized_method) }
diff --git a/test/memoist_test.rb b/test/memoist_test.rb
index 051cc9d..bf7c4d0 100644
--- a/test/memoist_test.rb
+++ b/test/memoist_test.rb
@@ -181,9 +181,37 @@ def counter
     memoize :counter
   end
 
+  class Book
+    extend Memoist
+    STATUSES = %w(new used)
+    CLASSIFICATION = %w(fiction nonfiction)
+    GENRES = %w(humor romance reference sci-fi classic philosophy)
+
+    attr_reader :title, :author
+    def initialize(title, author)
+      @title = title
+      @author = author
+    end
+
+    def full_title
+      "#{@title} by #{@author}"
+    end
+    memoize :full_title
+
+    class << self
+      extend Memoist
+
+      def all_types
+        STATUSES.product(CLASSIFICATION).product(GENRES).collect { |a| a.flatten }
+      end
+      memoize :all_types
+    end
+  end
+
   def setup
     @person = Person.new
     @calculator = Calculator.new
+    @book = Book.new("My Life", "Brian 'Fudge' Turmuck")
   end
 
   def test_memoization
@@ -256,6 +284,26 @@ def test_flush_cache
     assert_equal 2, @calculator.counter
   end
 
+  def test_class_flush_cache
+    @book.memoize_all
+    assert_equal "My Life by Brian 'Fudge' Turmuck", @book.full_title
+
+    Book.memoize_all
+    assert_instance_of Array, Book.instance_variable_get(:@_memoized_all_types)
+    Book.flush_cache
+    assert_equal false, Book.instance_variable_defined?(:@_memoized_all_types)
+  end
+
+  def test_class_flush_cache_preserves_instances
+    @book.memoize_all
+    Book.memoize_all
+    assert_equal "My Life by Brian 'Fudge' Turmuck", @book.full_title
+    
+    Book.flush_cache
+    assert_equal false, Book.instance_variable_defined?(:@_memoized_all_types)
+    assert_equal "My Life by Brian 'Fudge' Turmuck", @book.full_title
+  end
+
   def test_unmemoize_all
     assert_equal 1, @calculator.counter
 
@@ -342,6 +390,16 @@ def test_memoization_cache_is_different_for_each_instance
     assert_equal 1, Calculator.new.counter
   end
 
+  def test_memoization_class_variables
+    @book.memoize_all
+    assert_equal "My Life by Brian 'Fudge' Turmuck", @book.instance_variable_get(:@_memoized_full_title)
+    assert_equal "My Life by Brian 'Fudge' Turmuck", @book.full_title
+
+    Book.memoize_all
+    assert_instance_of Array, Book.instance_variable_get(:@_memoized_all_types)
+    assert_equal 24, Book.all_types.count
+  end
+
   def test_memoized_is_not_affected_by_freeze
     @person.freeze
     assert_equal "Josh", @person.name