[ruby/error_highlight] Show a dedicated snippet for "wrong number of arguments" error

This is an experimental implementation for
https://bugs.ruby-lang.org/issues/21543.

```
test.rb:2:in 'Object#foo': wrong number of arguments (given 1, expected 2) (ArgumentError)

    caller: test.rb:6
    |   foo(1)
        ^^^
    callee: test.rb:2
    | def foo(x, y)
          ^^^
        from test.rb:6:in 'Object#bar'
        from test.rb:10:in 'Object#baz'
        from test.rb:13:in '<main>'
```

https://github.com/ruby/error_highlight/commit/21e974e1c4

3 files changed, 346 insertions, 20 deletions

diff --git a/lib/error_highlight/base.rb b/lib/error_highlight/base.rb
index 14e0ce5785..b4a31f8e80 100644
--- a/lib/error_highlight/base.rb
+++ b/lib/error_highlight/base.rb
@@ -239,6 +239,20 @@ module ErrorHighlight
       when :OP_CDECL
         spot_op_cdecl
 
+      when :DEFN
+        raise NotImplementedError if @point_type != :name
+        spot_defn
+
+      when :DEFS
+        raise NotImplementedError if @point_type != :name
+        spot_defs
+
+      when :LAMBDA
+        spot_lambda
+
+      when :ITER
+        spot_iter
+
       when :call_node
         case @point_type
         when :name
@@ -280,6 +294,30 @@ module ErrorHighlight
       when :constant_path_operator_write_node
         prism_spot_constant_path_operator_write
 
+      when :def_node
+        case @point_type
+        when :name
+          prism_spot_def_for_name
+        when :args
+          raise NotImplementedError
+        end
+
+      when :lambda_node
+        case @point_type
+        when :name
+          prism_spot_lambda_for_name
+        when :args
+          raise NotImplementedError
+        end
+
+      when :block_node
+        case @point_type
+        when :name
+          prism_spot_block_for_name
+        when :args
+          raise NotImplementedError
+        end
+
       end
 
       if @snippet && @beg_column && @end_column && @beg_column < @end_column
@@ -621,6 +659,55 @@ module ErrorHighlight
       end
     end
 
+    # Example:
+    #   def bar; end
+    #       ^^^
+    def spot_defn
+      mid, = @node.children
+      fetch_line(@node.first_lineno)
+      if @snippet.match(/\Gdef\s+(#{ Regexp.quote(mid) }\b)/, @node.first_column)
+        @beg_column = $~.begin(1)
+        @end_column = $~.end(1)
+      end
+    end
+
+    # Example:
+    #   def Foo.bar; end
+    #          ^^^^
+    def spot_defs
+      nd_recv, mid, = @node.children
+      fetch_line(nd_recv.last_lineno)
+      if @snippet.match(/\G\s*(\.\s*#{ Regexp.quote(mid) }\b)/, nd_recv.last_column)
+        @beg_column = $~.begin(1)
+        @end_column = $~.end(1)
+      end
+    end
+
+    # Example:
+    #   -> { ... }
+    #   ^^
+    def spot_lambda
+      fetch_line(@node.first_lineno)
+      if @snippet.match(/\G->/, @node.first_column)
+        @beg_column = $~.begin(0)
+        @end_column = $~.end(0)
+      end
+    end
+
+    # Example:
+    #   lambda { ... }
+    #          ^
+    #   define_method :foo do
+    #                      ^^
+    def spot_iter
+      _nd_fcall, nd_scope = @node.children
+      fetch_line(nd_scope.first_lineno)
+      if @snippet.match(/\G(?:do\b|\{)/, nd_scope.first_column)
+        @beg_column = $~.begin(0)
+        @end_column = $~.end(0)
+      end
+    end
+
     def fetch_line(lineno)
       @beg_lineno = @end_lineno = lineno
       @snippet = @fetch[lineno]
@@ -826,6 +913,31 @@ module ErrorHighlight
         prism_location(@node.binary_operator_loc.chop)
       end
     end
+
+    # Example:
+    #   def foo()
+    #       ^^^
+    def prism_spot_def_for_name
+      location = @node.name_loc
+      location = location.join(@node.operator_loc) if @node.operator_loc
+      prism_location(location)
+    end
+
+    # Example:
+    #   -> x, y { }
+    #   ^^
+    def prism_spot_lambda_for_name
+      prism_location(@node.operator_loc)
+    end
+
+    # Example:
+    #   lambda { }
+    #          ^
+    #   define_method :foo do |x, y|
+    #                      ^
+    def prism_spot_block_for_name
+      prism_location(@node.opening_loc)
+    end
   end
 
   private_constant :Spotter
diff --git a/lib/error_highlight/core_ext.rb b/lib/error_highlight/core_ext.rb
index b69093f74e..2fb07f2e65 100644
--- a/lib/error_highlight/core_ext.rb
+++ b/lib/error_highlight/core_ext.rb
@@ -3,9 +3,38 @@ require_relative "formatter"
 module ErrorHighlight
   module CoreExt
     private def generate_snippet
-      spot = ErrorHighlight.spot(self)
-      return "" unless spot
-      return ErrorHighlight.formatter.message_for(spot)
+      if ArgumentError === self && message =~ /\A(?:wrong number of arguments|missing keyword|unknown keyword|no keywords accepted)\b/
+        locs = self.backtrace_locations
+        return "" if locs.size < 2
+        callee_loc, caller_loc = locs
+        callee_spot = ErrorHighlight.spot(self, backtrace_location: callee_loc, point_type: :name)
+        caller_spot = ErrorHighlight.spot(self, backtrace_location: caller_loc, point_type: :name)
+        if caller_spot && callee_spot &&
+            caller_loc.path == callee_loc.path &&
+            caller_loc.lineno == callee_loc.lineno &&
+            caller_spot == callee_spot
+          callee_loc = callee_spot = nil
+        end
+        ret = +"\n"
+        [["caller", caller_loc, caller_spot], ["callee", callee_loc, callee_spot]].each do |header, loc, spot|
+          out = nil
+          if loc
+            out = "    #{ header }: #{ loc.path }:#{ loc.lineno }"
+            if spot
+              _, _, snippet, highlight = ErrorHighlight.formatter.message_for(spot).lines
+              out += "\n    | #{ snippet }      #{ highlight }"
+            else
+              out += "\n    (cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+            end
+          end
+          ret << "\n" + out if out
+        end
+        ret
+      else
+        spot = ErrorHighlight.spot(self)
+        return "" unless spot
+        return ErrorHighlight.formatter.message_for(spot)
+      end
     end
 
     if Exception.method_defined?(:detailed_message)
diff --git a/test/error_highlight/test_error_highlight.rb b/test/error_highlight/test_error_highlight.rb
index 8aa5eb9c8d..09b0579f8a 100644
--- a/test/error_highlight/test_error_highlight.rb
+++ b/test/error_highlight/test_error_highlight.rb
@@ -44,14 +44,16 @@ class ErrorHighlightTest < Test::Unit::TestCase
     def assert_error_message(klass, expected_msg, &blk)
       omit unless klass < ErrorHighlight::CoreExt
       err = assert_raise(klass, &blk)
-      spot = ErrorHighlight.spot(err)
-      if spot
-        assert_kind_of(Integer, spot[:first_lineno])
-        assert_kind_of(Integer, spot[:first_column])
-        assert_kind_of(Integer, spot[:last_lineno])
-        assert_kind_of(Integer, spot[:last_column])
-        assert_kind_of(String, spot[:snippet])
-        assert_kind_of(Array, spot[:script_lines])
+      unless klass == ArgumentError && err.message =~ /\A(?:wrong number of arguments|missing keyword|unknown keyword|no keywords accepted)\b/
+        spot = ErrorHighlight.spot(err)
+        if spot
+          assert_kind_of(Integer, spot[:first_lineno])
+          assert_kind_of(Integer, spot[:first_column])
+          assert_kind_of(Integer, spot[:last_lineno])
+          assert_kind_of(Integer, spot[:last_column])
+          assert_kind_of(String, spot[:snippet])
+          assert_kind_of(Array, spot[:script_lines])
+        end
       end
       assert_equal(preprocess(expected_msg).chomp, err.detailed_message(highlight: false).sub(/ \((?:NoMethod|Name)Error\)/, ""))
     end
@@ -1111,12 +1113,13 @@ no implicit conversion from nil to integer (TypeError)
   end
 
   def test_args_ATTRASGN_1
-    v = []
-    assert_error_message(ArgumentError, <<~END) do
-wrong number of arguments (given 1, expected 2..3) (ArgumentError)
+    v = method(:raise).to_proc
+    recv = NEW_MESSAGE_FORMAT ? "an instance of Proc" : v.inspect
+    assert_error_message(NoMethodError, <<~END) do
+undefined method `[]=' for #{ recv }
 
       v [ ] = 1
-         ^^^^^^
+        ^^^^^
     END
 
       v [ ] = 1
@@ -1199,16 +1202,16 @@ no implicit conversion from nil to integer (TypeError)
   end
 
   def test_args_OP_ASGN1_aref_2
-    v = []
+    v = method(:raise).to_proc
 
     assert_error_message(ArgumentError, <<~END) do
-wrong number of arguments (given 0, expected 1..2) (ArgumentError)
+ArgumentError (ArgumentError)
 
-      v [ ] += 42
-         ^^^^^^^^
+      v [ArgumentError] += 42
+         ^^^^^^^^^^^^^^^^^^^^
     END
 
-      v [ ] += 42
+      v [ArgumentError] += 42
     end
   end
 
@@ -1453,6 +1456,188 @@ undefined method `foo' for #{ NIL_RECV_MESSAGE }
     end
   end
 
+  begin
+    ->{}.call(1)
+  rescue ArgumentError => exc
+    MethodDefLocationSupported =
+      RubyVM::AbstractSyntaxTree.respond_to?(:node_id_for_backtrace_location) &&
+      RubyVM::AbstractSyntaxTree.node_id_for_backtrace_location(exc.backtrace_locations.first)
+  end
+
+  WRONG_NUMBER_OF_ARGUMENTS_LIENO = __LINE__ + 1
+  def wrong_number_of_arguments_test(x, y)
+    x + y
+  end
+
+  def test_wrong_number_of_arguments_for_method
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+wrong number of arguments (given 1, expected 2) (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       wrong_number_of_arguments_test(1)
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    callee: #{ __FILE__ }:#{ WRONG_NUMBER_OF_ARGUMENTS_LIENO }
+    #{
+      MethodDefLocationSupported ?
+   "|   def wrong_number_of_arguments_test(x, y)
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      wrong_number_of_arguments_test(1)
+    end
+  end
+
+  KEYWORD_TEST_LINENO = __LINE__ + 1
+  def keyword_test(kw1:, kw2:, kw3:)
+    kw1 + kw2 + kw3
+  end
+
+  def test_missing_keyword
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+missing keyword: :kw3 (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       keyword_test(kw1: 1, kw2: 2)
+            ^^^^^^^^^^^^
+    callee: #{ __FILE__ }:#{ KEYWORD_TEST_LINENO }
+    #{
+      MethodDefLocationSupported ?
+   "|   def keyword_test(kw1:, kw2:, kw3:)
+            ^^^^^^^^^^^^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      keyword_test(kw1: 1, kw2: 2)
+    end
+  end
+
+  def test_unknown_keyword
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+unknown keyword: :kw4 (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       keyword_test(kw1: 1, kw2: 2, kw3: 3, kw4: 4)
+            ^^^^^^^^^^^^
+    callee: #{ __FILE__ }:#{ KEYWORD_TEST_LINENO }
+    #{
+      MethodDefLocationSupported ?
+   "|   def keyword_test(kw1:, kw2:, kw3:)
+            ^^^^^^^^^^^^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      keyword_test(kw1: 1, kw2: 2, kw3: 3, kw4: 4)
+    end
+  end
+
+  WRONG_NUBMER_OF_ARGUMENTS_TEST2_LINENO = __LINE__ + 1
+  def wrong_number_of_arguments_test2(
+    long_argument_name_x,
+    long_argument_name_y,
+    long_argument_name_z
+  )
+    long_argument_name_x + long_argument_name_y + long_argument_name_z
+  end
+
+  def test_wrong_number_of_arguments_for_method2
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+wrong number of arguments (given 1, expected 3) (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       wrong_number_of_arguments_test2(1)
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    callee: #{ __FILE__ }:#{ WRONG_NUBMER_OF_ARGUMENTS_TEST2_LINENO }
+    #{
+      MethodDefLocationSupported ?
+   "|   def wrong_number_of_arguments_test2(
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      wrong_number_of_arguments_test2(1)
+    end
+  end
+
+  def test_wrong_number_of_arguments_for_lambda_literal
+    v = -> {}
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+wrong number of arguments (given 1, expected 0) (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       v.call(1)
+             ^^^^^
+    callee: #{ __FILE__ }:#{ lineno - 1 }
+    #{
+      MethodDefLocationSupported ?
+   "|     v = -> {}
+              ^^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      v.call(1)
+    end
+  end
+
+  def test_wrong_number_of_arguments_for_lambda_method
+    v = lambda { }
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+wrong number of arguments (given 1, expected 0) (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       v.call(1)
+             ^^^^^
+    callee: #{ __FILE__ }:#{ lineno - 1 }
+    #{
+      MethodDefLocationSupported ?
+   "|     v = lambda { }
+                     ^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      v.call(1)
+    end
+  end
+
+  DEFINE_METHOD_TEST_LINENO = __LINE__ + 1
+  define_method :define_method_test do |x, y|
+    x + y
+  end
+
+  def test_wrong_number_of_arguments_for_define_method
+    v = lambda { }
+    lineno = __LINE__
+    assert_error_message(ArgumentError, <<~END) do
+wrong number of arguments (given 1, expected 2) (ArgumentError)
+
+    caller: #{ __FILE__ }:#{ lineno + 16 }
+    |       define_method_test(1)
+            ^^^^^^^^^^^^^^^^^^
+    callee: #{ __FILE__ }:#{ DEFINE_METHOD_TEST_LINENO }
+    #{
+      MethodDefLocationSupported ?
+   "|   define_method :define_method_test do |x, y|
+                                          ^^" :
+   "(cannot create a snippet of the method definition; use Ruby 3.5 or later)"
+   }
+    END
+
+      define_method_test(1)
+    end
+  end
+
   def test_spoofed_filename
     Tempfile.create(["error_highlight_test", ".rb"], binmode: true) do |tmp|
       tmp << "module Dummy\nend\n"