From 8c53b279da67d10d7d2befd5a0762c08ec4cf8a1 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

---
 lib/memo_wise.rb       | 20 ++++++++++++++++++++
 spec/memo_wise_spec.rb | 17 +++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/lib/memo_wise.rb b/lib/memo_wise.rb
index dace3d6..5f61e36 100644
--- a/lib/memo_wise.rb
+++ b/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
 
diff --git a/spec/memo_wise_spec.rb b/spec/memo_wise_spec.rb
index 2a3f4b9..0fca4ee 100644
--- a/spec/memo_wise_spec.rb
+++ b/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