Split out *args and **kwargs cases

lib/memo_wise.rb

@@ -8,13 +8,13 @@ def initialize(*values)
     super
   end
 
-  def self.prepended(target)
+  def self.prepended(target) # rubocop:disable Metrics/PerceivedComplexity
     class << target
       # Implements memoization for the given method name.
       #
       # @param method_name [Symbol]
       #   Name of method for which to implement memoization.
-      def memo_wise(method_name)
+      def memo_wise(method_name) # rubocop:disable Metrics/PerceivedComplexity
         unless method_name.is_a?(Symbol)
           raise ArgumentError,
                 "#{method_name.inspect} must be a Symbol"
@@ -32,7 +32,11 @@ def memo_wise(method_name)
         alias_method not_memoized_name, method_name
         private not_memoized_name
 
-        if instance_method(method_name).arity.zero?
+        method = instance_method(method_name)
+
+        # Zero-arg methods can use simpler/more performant logic because the
+        # hash key is just the method name.
+        if method.arity.zero?
           module_eval <<-END_OF_METHOD, __FILE__, __LINE__ + 1
             def #{method_name}
               @_memo_wise.fetch(:#{method_name}) do
@@ -41,14 +45,36 @@ def #{method_name}
             end
           END_OF_METHOD
         else
+          # If our method has arguments, we need to separate out our handling of
+          # normal args vs. keyword args due to the changes in Ruby 3.
+          # See: <link>
+          # By only including logic for *args or **kwargs when they are used in
+          # the method, we can avoid allocating unnecessary arrays and hashes.
+          has_arg = method.parameters.any? do |(param, _)|
+            param == :req || param == :opt || param == :rest # rubocop:disable Style/MultipleComparison
+          end
+          has_kwarg = method.parameters.any? do |(param, _)|
+            param == :keyreq || param == :key || param == :keyrest # rubocop:disable Style/MultipleComparison
+          end
+
+          if has_arg && has_kwarg
+            args_str = "(*args, **kwargs)"
+            fetch_key = "[args, kwargs].freeze"
+          elsif has_arg
+            args_str = "(*args)"
+            fetch_key = "args"
+          else
+            args_str = "(**kwargs)"
+            fetch_key = "kwargs"
+          end
+
           # Note that we don't need to freeze args before using it as a hash key
           # because Ruby always copies argument arrays when splatted.
           module_eval <<-END_OF_METHOD, __FILE__, __LINE__ + 1
-            def #{method_name}(*args, **kwargs)
+            def #{method_name}#{args_str}
               hash = @_memo_wise[:#{method_name}]
-              key = [args,kwargs]
-              hash.fetch(key) do
-                hash[key] = #{not_memoized_name}(*args, **kwargs)
+              hash.fetch(#{fetch_key}) do
+                hash[#{fetch_key}] = #{not_memoized_name}#{args_str}
               end
             end
           END_OF_METHOD
@@ -61,7 +87,7 @@ def #{method_name}(*args, **kwargs)
     end
   end
 
-  def reset_memo_wise(method_name, *args, **kwargs)
+  def reset_memo_wise(method_name, *args, **kwargs) # rubocop:disable Metrics/PerceivedComplexity
     unless method_name.is_a?(Symbol)
       raise ArgumentError, "#{method_name.inspect} must be a Symbol"
     end
@@ -73,7 +99,22 @@ def reset_memo_wise(method_name, *args, **kwargs)
     if args.empty? && kwargs.empty?
       @_memo_wise.delete(method_name)
     else
-      @_memo_wise[method_name].delete([args, kwargs])
+      method = self.class.instance_method(method_name)
+
+      has_arg = method.parameters.any? do |(param, _)|
+        param == :req || param == :opt || param == :rest # rubocop:disable Style/MultipleComparison
+      end
+      has_kwarg = method.parameters.any? do |(param, _)|
+        param == :keyreq || param == :key || param == :keyrest # rubocop:disable Style/MultipleComparison
+      end
+
+      if has_arg && has_kwarg
+        @_memo_wise[method_name].delete([args, kwargs])
+      elsif has_arg
+        @_memo_wise[method_name].delete(args)
+      else
+        @_memo_wise[method_name].delete(kwargs)
+      end
     end
   end