always set memo_wise hash if not set for some reason

panorama-ed · Oct 31, 2023 · fae1c66 · fae1c66
1 parent 9b289a9
commit fae1c66
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/lib/memo_wise.rb b/lib/memo_wise.rb
@@ -181,6 +181,7 @@ def inherited(subclass)
         when MemoWise::InternalAPI::NONE
           klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}
+              @_memo_wise ||= {}
               @_memo_wise.fetch(:#{method_name}) do
                 @_memo_wise[:#{method_name}] = #{original_memo_wised_name}
               end
@@ -190,6 +191,7 @@ def #{method_name}
           key = method.parameters.first.last
           klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
+              @_memo_wise ||= {}
               _memo_wise_hash = (@_memo_wise[:#{method_name}] ||= {})
               _memo_wise_hash.fetch(#{key}) do
                 _memo_wise_hash[#{key}] = #{original_memo_wised_name}(#{MemoWise::InternalAPI.call_str(method)})
@@ -211,6 +213,7 @@ def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
           end
           klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
+              @_memo_wise ||= {}
               _memo_wise_hash = (@_memo_wise[:#{method_name}] ||= {})
               #{layers.join("\n  ")}
             end
@@ -220,6 +223,7 @@ def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
           # { method_name: { args => { kwargs => memoized_value } } }
           klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}(*args, **kwargs)
+              @_memo_wise ||= {}
               _memo_wise_hash = (@_memo_wise[:#{method_name}] ||= {})
               _memo_wise_kwargs_hash = _memo_wise_hash.fetch(args) do
                 _memo_wise_hash[args] = {}
@@ -232,6 +236,7 @@ def #{method_name}(*args, **kwargs)
         else # MemoWise::InternalAPI::SPLAT, MemoWise::InternalAPI::DOUBLE_SPLAT
           klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
+              @_memo_wise ||= {}
               _memo_wise_hash = (@_memo_wise[:#{method_name}] ||= {})
               _memo_wise_key = #{MemoWise::InternalAPI.key_str(method)}
               _memo_wise_hash.fetch(_memo_wise_key) do
@@ -552,6 +557,7 @@ def reset_memo_wise(method_name = nil, *args, **kwargs)
       raise ArgumentError, "Provided args when method_name = nil" unless args.empty?
       raise ArgumentError, "Provided kwargs when method_name = nil" unless kwargs.empty?
 
+      @_memo_wise ||= {}
       @_memo_wise.clear
       return
     end

diff --git 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