lib/memo_wise.rb in memo_wise-1.2.0 vs lib/memo_wise.rb in memo_wise-1.3.0

@@ -54,11 +54,11 @@
   # for more information.
   #
   # :nocov:
   all_args = RUBY_VERSION < "2.7" ? "*" : "..."
   # :nocov:
-  class_eval <<-END_OF_METHOD, __FILE__, __LINE__ + 1
+  class_eval <<~HEREDOC, __FILE__, __LINE__ + 1
     # On Ruby 2.7 or greater:
     #
     # def initialize(...)
     #   MemoWise::InternalAPI.create_memo_wise_state!(self)
     #   super
@@ -73,11 +73,11 @@
 
     def initialize(#{all_args})
       MemoWise::InternalAPI.create_memo_wise_state!(self)
       super
     end
-  END_OF_METHOD
+  HEREDOC
 
   # @private
   #
   # Private setup method, called automatically by `prepend MemoWise` in a class.
   #
@@ -154,10 +154,21 @@
           # {Class#singleton_class}.
           # See: https://medium.com/@leo_hetsch/demystifying-singleton-classes-in-ruby-caf3fa4c9d91
           klass = klass.singleton_class
         end
 
+        if klass.singleton_class?
+          # This ensures that a memoized method defined on a parent class can
+          # still be used in a child class.
+          klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
+            def inherited(subclass)
+              super
+              MemoWise::InternalAPI.create_memo_wise_state!(subclass)
+            end
+          HEREDOC
+        end
+
         raise ArgumentError, "#{method_name.inspect} must be a Symbol" unless method_name.is_a?(Symbol)
 
         api = MemoWise::InternalAPI.new(klass)
         visibility = api.method_visibility(method_name)
         original_memo_wised_name = MemoWise::InternalAPI.original_memo_wised_name(method_name)
@@ -165,53 +176,41 @@
 
         klass.send(:alias_method, original_memo_wised_name, method_name)
         klass.send(:private, original_memo_wised_name)
 
         method_arguments = MemoWise::InternalAPI.method_arguments(method)
-        # `@_memo_wise_indices` stores the `@_memo_wise` indices of different
-        # method names. We only use this data structure when resetting or
-        # presetting memoization. It looks like:
-        #   {
-        #     single_arg_method_name: 0,
-        #     other_single_arg_method_name: 1
-        #   }
-        memo_wise_indices = klass.instance_variable_get(:@_memo_wise_indices)
-        memo_wise_indices ||= klass.instance_variable_set(:@_memo_wise_indices, {})
-        index = klass.instance_variable_get(:@_memo_wise_index_counter) || 0
+        index = MemoWise::InternalAPI.next_index!(klass, method_name)
 
-        memo_wise_indices[method_name] = index
-        klass.instance_variable_set(:@_memo_wise_index_counter, index + 1)
-
         case method_arguments
         when MemoWise::InternalAPI::NONE
           # Zero-arg methods can use simpler/more performant logic because the
           # hash key is just the method name.
-          klass.module_eval <<-END_OF_METHOD, __FILE__, __LINE__ + 1
+          klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}
-              _memo_wise_output = @_memo_wise[#{index}]
-              if _memo_wise_output || @_memo_wise_sentinels[#{index}]
-                _memo_wise_output
+              if @_memo_wise_sentinels[#{index}]
+                @_memo_wise[#{index}]
               else
+                ret = @_memo_wise[#{index}] = #{original_memo_wised_name}
                 @_memo_wise_sentinels[#{index}] = true
-                @_memo_wise[#{index}] = #{original_memo_wised_name}
+                ret
               end
             end
-          END_OF_METHOD
+          HEREDOC
         when MemoWise::InternalAPI::ONE_REQUIRED_POSITIONAL, MemoWise::InternalAPI::ONE_REQUIRED_KEYWORD
           key = method.parameters.first.last
 
-          klass.module_eval <<-END_OF_METHOD, __FILE__, __LINE__ + 1
+          klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
               _memo_wise_hash = (@_memo_wise[#{index}] ||= {})
               _memo_wise_output = _memo_wise_hash[#{key}]
               if _memo_wise_output || _memo_wise_hash.key?(#{key})
                 _memo_wise_output
               else
                 _memo_wise_hash[#{key}] = #{original_memo_wised_name}(#{MemoWise::InternalAPI.call_str(method)})
               end
             end
-          END_OF_METHOD
+          HEREDOC
         # MemoWise::InternalAPI::MULTIPLE_REQUIRED, MemoWise::InternalAPI::SPLAT,
         # MemoWise::InternalAPI::DOUBLE_SPLAT, MemoWise::InternalAPI::SPLAT_AND_DOUBLE_SPLAT
         else
           # NOTE: When benchmarking this implementation against something like:
           #
@@ -223,21 +222,21 @@
           # is because this case uses a more complex hash key (see
           # `MemoWise::InternalAPI.key_str`), and hashing that key has less
           # consistent performance. In general, this should still be faster for
           # truthy results because `Hash#[]` generally performs hash lookups
           # faster than `Hash#fetch`.
-          klass.module_eval <<-END_OF_METHOD, __FILE__, __LINE__ + 1
+          klass.module_eval <<~HEREDOC, __FILE__, __LINE__ + 1
             def #{method_name}(#{MemoWise::InternalAPI.args_str(method)})
               _memo_wise_hash = (@_memo_wise[#{index}] ||= {})
               _memo_wise_key = #{MemoWise::InternalAPI.key_str(method)}
               _memo_wise_output = _memo_wise_hash[_memo_wise_key]
               if _memo_wise_output || _memo_wise_hash.key?(_memo_wise_key)
                 _memo_wise_output
               else
                 _memo_wise_hash[_memo_wise_key] = #{original_memo_wised_name}(#{MemoWise::InternalAPI.call_str(method)})
               end
             end
-          END_OF_METHOD
+          HEREDOC
         end
 
         klass.send(visibility, method_name)
       end
     end