From 21040dc578ce76db3cbb2ed9577d2b11039b7d8c Mon Sep 17 00:00:00 2001
From: Jacob Evelyn <jevelyn@panoramaed.com>
Date: Tue, 30 Jan 2024 09:58:06 -0500
Subject: [PATCH] Prepend ivar setup in included modules

Fixes #302

Co-authored-by: alpaca-tc <alpaca-tc@alpaca.tc>
---
 lib/memo_wise.rb       | 23 +++++++++++++++++++----
 spec/memo_wise_spec.rb | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/lib/memo_wise.rb b/lib/memo_wise.rb
index dace3d6..795dc06 100644
--- a/lib/memo_wise.rb
+++ b/lib/memo_wise.rb
@@ -54,10 +54,9 @@ module MemoWise
   # [this article](https://www.ruby-lang.org/en/news/2019/12/12/separation-of-positional-and-keyword-arguments-in-ruby-3-0/)
   # for more information.
   #
+
   # :nocov:
-  all_args = RUBY_VERSION < "2.7" ? "*" : "..."
-  # :nocov:
-  class_eval <<~HEREDOC, __FILE__, __LINE__ + 1
+  INITIALIZE_LITERAL = <<~HEREDOC
     # On Ruby 2.7 or greater:
     #
     # def initialize(...)
@@ -72,11 +71,14 @@ module MemoWise
     #   super
     # end
 
-    def initialize(#{all_args})
+    def initialize(#{RUBY_VERSION < "2.7" ? "*" : "..."})
       MemoWise::InternalAPI.create_memo_wise_state!(self)
       super
     end
   HEREDOC
+  # :nocov:
+
+  class_eval(INITIALIZE_LITERAL, __FILE__, __LINE__ + 1)
 
   module CreateMemoWiseStateOnExtended
     def extended(base)
@@ -94,6 +96,15 @@ def inherited(subclass)
   end
   private_constant(:CreateMemoWiseStateOnInherited)
 
+  module CreateMemoWiseStateOnIncluded
+    def included(base)
+      base.prepend(Module.new do
+        class_eval(INITIALIZE_LITERAL, __FILE__, __LINE__ + 1)
+      end)
+    end
+  end
+  private_constant(:CreateMemoWiseStateOnIncluded)
+
   # @private
   #
   # Private setup method, called automatically by `prepend MemoWise` in a class.
@@ -176,6 +187,10 @@ def memo_wise(method_name_or_hash)
         if klass.is_a?(Class) && !klass.singleton_class?
           klass.singleton_class.prepend(CreateMemoWiseStateOnInherited)
         else
+          if klass.is_a?(Module) && !klass.singleton_class?
+            klass.singleton_class.prepend(CreateMemoWiseStateOnIncluded)
+          end
+
           klass.prepend(CreateMemoWiseStateOnInherited)
         end
 
diff --git a/spec/memo_wise_spec.rb b/spec/memo_wise_spec.rb
index 2a3f4b9..5efaccc 100644
--- a/spec/memo_wise_spec.rb
+++ b/spec/memo_wise_spec.rb
@@ -351,11 +351,32 @@ def module2_method
           end
         end
 
+        let(:klass_with_initializer) do
+          Class.new do
+            include Module1
+            def initialize(*); end
+          end
+        end
+
+        let(:module_with_initializer) do
+          Module.new do
+            include Module1
+            def initialize(*); end
+          end
+        end
+
+        let(:klass_with_module_with_initializer) do
+          Class.new do
+            include Module3
+          end
+        end
+
         let(:instance) { klass.new }
 
         before(:each) do
           stub_const("Module1", module1)
           stub_const("Module2", module2)
+          stub_const("Module3", module_with_initializer)
         end
 
         it "memoizes inherited methods separately" do
@@ -364,6 +385,27 @@ def module2_method
           expect(Array.new(4) { instance.module2_method }).to all eq("module2_method")
           expect(instance.module2_method_counter).to eq(1)
         end
+
+        it "can memoize klass with initializer" do
+          instance = klass_with_initializer.new(true)
+          expect { instance.module1_method }.not_to raise_error
+
+          expect(Array.new(4) { instance.module1_method }).to all eq("module1_method")
+          expect(instance.module1_method_counter).to eq(1)
+        end
+
+        it "can memoize klass with initializer" do
+          instance = klass_with_module_with_initializer.new(true)
+          expect { instance.module1_method }.not_to raise_error
+
+          expect(Array.new(4) { instance.module1_method }).to all eq("module1_method")
+          expect(instance.module1_method_counter).to eq(1)
+        end
+
+        it "can reset klass with initializer" do
+          instance = klass_with_initializer.new(true)
+          expect { instance.reset_memo_wise }.not_to raise_error
+        end
       end
 
       context "when the class, its superclass, and its module all memoize methods" do