Prepend ivar setup in included modules

lib/memo_wise.rb

@@ -94,6 +94,20 @@ def inherited(subclass)
   end
   private_constant(:CreateMemoWiseStateOnInherited)
 
+  module CreateMemoWiseStateOnIncluded
+    def included(base)
+      return unless base.is_a?(Class) && !base.singleton_class?
+
+      base.prepend(Module.new do
+        def initialize(...)
+          MemoWise::InternalAPI.create_memo_wise_state!(self)
+          super
+        end
+      end)
+    end
+  end
+  private_constant(:CreateMemoWiseStateOnIncluded)
+
   # @private
   #
   # Private setup method, called automatically by `prepend MemoWise` in a class.
@@ -176,6 +190,12 @@ 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)
+          elsif klass.is_a?(Module)
+            klass.prepend(CreateMemoWiseStateOnIncluded)
+          end
+
           klass.prepend(CreateMemoWiseStateOnInherited)
         end
 

spec/memo_wise_spec.rb

@@ -351,6 +351,13 @@ def module2_method
           end
         end
 
+        let(:klass_with_initializer) do
+          Class.new do
+            include Module1
+            def initialize(*); end
+          end
+        end
+
         let(:instance) { klass.new }
 
         before(:each) do
@@ -364,6 +371,16 @@ 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
+        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