Prepend ivar setup in included modules

Fixes #302

Co-authored-by: alpaca-tc <[email protected]>

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
 

spec/adding_methods_spec.rb

@@ -60,7 +60,7 @@ def self.example; end
           end
         end
 
-        let(:expected_public_class_methods) { super() << :inherited }
+        let(:expected_public_class_methods) { super() + %i[inherited included] }
 
         it "adds expected public *instance* methods only" do
           expect { subject }.

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