From 89b0513a4fa0dc61aee9f821f231b926b3d8ce9a Mon Sep 17 00:00:00 2001
From: Nikita Vasilevsky <nvasilevskii@gmail.com>
Date: Tue, 1 Feb 2022 19:50:39 +0000
Subject: [PATCH] Support ruby 3.1 :args_forward shape change

---
 lib/yard.rb                              |  3 +++
 lib/yard/handlers/ruby/method_handler.rb |  2 +-
 lib/yard/parser/ruby/ast_node.rb         |  7 +++++++
 spec/handlers/method_handler_spec.rb     | 24 ++++++++++++++++++++++++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/lib/yard.rb b/lib/yard.rb
index 6afd7a0a9..3e14c4140 100644
--- a/lib/yard.rb
+++ b/lib/yard.rb
@@ -51,6 +51,9 @@ def self.ruby2?; @ruby2 ||= (RUBY_VERSION >= '2.0.0') end
 
   # @return [Boolean] whether YARD is being run in Ruby 3.0
   def self.ruby3?; @ruby3 ||= (RUBY_VERSION >= '3.0.0') end
+
+  # @return [Boolean] whether YARD is being run in Ruby 3.1
+  def self.ruby31?; @ruby31 ||= (RUBY_VERSION >= '3.1.0') end
 end
 
 # Keep track of Ruby version for compatibility code
diff --git a/lib/yard/handlers/ruby/method_handler.rb b/lib/yard/handlers/ruby/method_handler.rb
index 5030a9658..c779c5b39 100644
--- a/lib/yard/handlers/ruby/method_handler.rb
+++ b/lib/yard/handlers/ruby/method_handler.rb
@@ -97,7 +97,7 @@ def format_args
       params << ['**' + args.double_splat_param.source, nil]
     end
 
-    params << ['&' + args.block_param.source, nil] if args.block_param
+    params << ['&' + args.block_param.source, nil] if args.block_param && !args.args_forward
 
     params
   end
diff --git a/lib/yard/parser/ruby/ast_node.rb b/lib/yard/parser/ruby/ast_node.rb
index 0d19e115d..c9deb3ed1 100644
--- a/lib/yard/parser/ruby/ast_node.rb
+++ b/lib/yard/parser/ruby/ast_node.rb
@@ -426,6 +426,13 @@ def double_splat_param
         def block_param
           self[-1] ? self[-1][0] : nil
         end
+
+        def args_forward
+          # shape is (required, optional, rest, more, keyword, keyword_rest, block)
+          # Ruby 3.1 moves :args_forward from rest to keyword_rest
+          args_index = YARD.ruby31? ? -2 : 2
+          self[args_index].type == :args_forward if self[args_index]
+        end
       end
 
       class MethodCallNode < AstNode
diff --git a/spec/handlers/method_handler_spec.rb b/spec/handlers/method_handler_spec.rb
index 6cfb1cc6c..5cc2ff69a 100644
--- a/spec/handlers/method_handler_spec.rb
+++ b/spec/handlers/method_handler_spec.rb
@@ -72,6 +72,30 @@ def endless = true
     expect(P('Bar#endless').signature).to eq "def endless"
   end if YARD.ruby3?
 
+  it "handles method with arguments forwarding" do
+    YARD.parse_string <<-EOF
+      class Bar
+        def method_with_forwarding(...)
+          forward_to(...)
+        end
+      end
+    EOF
+
+    expect(P('Bar#method_with_forwarding').signature).to eq "def method_with_forwarding(...)"
+  end
+
+  it "handles method with anonymous block" do
+    YARD.parse_string <<-EOF
+      class Bar
+        def anon_block_method(&)
+          baz(a, &)
+        end
+      end
+    EOF
+
+    expect(P('Bar#anon_block_method').signature).to eq "def anon_block_method(&)"
+  end if YARD.ruby31?
+
   it "handles endless method definitions with parameters" do
     YARD.parse_string <<-EOF
       class Bar