From 13738064095ac3220f8237cdda2b20884c5cf184 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Fri, 26 Apr 2024 14:05:13 -0400
Subject: [PATCH] Support symbol procs in definition

---
 lib/ruby_lsp/listeners/definition.rb          | 23 +++++++++++++------
 lib/ruby_lsp/requests/definition.rb           |  2 +-
 test/requests/definition_expectations_test.rb | 10 ++++----
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/lib/ruby_lsp/listeners/definition.rb b/lib/ruby_lsp/listeners/definition.rb
index fefc3f020..0757b5e2b 100644
--- a/lib/ruby_lsp/listeners/definition.rb
+++ b/lib/ruby_lsp/listeners/definition.rb
@@ -30,6 +30,7 @@ def initialize(response_builder, global_state, uri, nesting, dispatcher, typeche
         dispatcher.register(
           self,
           :on_call_node_enter,
+          :on_block_argument_node_enter,
           :on_constant_read_node_enter,
           :on_constant_path_node_enter,
         )
@@ -42,10 +43,21 @@ def on_call_node_enter(node)
         if message == :require || message == :require_relative
           handle_require_definition(node)
         else
-          handle_method_definition(node)
+          handle_method_definition(message.to_s, self_receiver?(node))
         end
       end
 
+      sig { params(node: Prism::BlockArgumentNode).void }
+      def on_block_argument_node_enter(node)
+        expression = node.expression
+        return unless expression.is_a?(Prism::SymbolNode)
+
+        value = expression.value
+        return unless value
+
+        handle_method_definition(value, false)
+      end
+
       sig { params(node: Prism::ConstantPathNode).void }
       def on_constant_path_node_enter(node)
         name = constant_name(node)
@@ -64,12 +76,9 @@ def on_constant_read_node_enter(node)
 
       private
 
-      sig { params(node: Prism::CallNode).void }
-      def handle_method_definition(node)
-        message = node.message
-        return unless message
-
-        methods = if self_receiver?(node)
+      sig { params(message: String, self_receiver: T::Boolean).void }
+      def handle_method_definition(message, self_receiver)
+        methods = if self_receiver
           @index.resolve_method(message, @nesting.join("::"))
         else
           # If the method doesn't have a receiver, then we provide a few candidates to jump to
diff --git a/lib/ruby_lsp/requests/definition.rb b/lib/ruby_lsp/requests/definition.rb
index a3b014a0b..a3de7f429 100644
--- a/lib/ruby_lsp/requests/definition.rb
+++ b/lib/ruby_lsp/requests/definition.rb
@@ -47,7 +47,7 @@ def initialize(document, global_state, position, dispatcher, typechecker_enabled
 
         target, parent, nesting = document.locate_node(
           position,
-          node_types: [Prism::CallNode, Prism::ConstantReadNode, Prism::ConstantPathNode],
+          node_types: [Prism::CallNode, Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::BlockArgumentNode],
         )
 
         if target.is_a?(Prism::ConstantReadNode) && parent.is_a?(Prism::ConstantPathNode)
diff --git a/test/requests/definition_expectations_test.rb b/test/requests/definition_expectations_test.rb
index 36ba75a34..d8ca6b987 100644
--- a/test/requests/definition_expectations_test.rb
+++ b/test/requests/definition_expectations_test.rb
@@ -407,6 +407,8 @@ def test_definition_precision_for_methods_with_block_arguments
     source = <<~RUBY
       class Foo
         def foo(&block); end
+
+        def argument; end
       end
 
       bar.foo(&:argument)
@@ -417,16 +419,16 @@ def foo(&block); end
       server.process_message(
         id: 1,
         method: "textDocument/definition",
-        params: { textDocument: { uri: uri }, position: { character: 12, line: 4 } },
+        params: { textDocument: { uri: uri }, position: { character: 12, line: 6 } },
       )
-      assert_empty(server.pop_response.response)
+      assert_equal(3, server.pop_response.response.first.range.start.line)
 
       server.process_message(
         id: 1,
         method: "textDocument/definition",
-        params: { textDocument: { uri: uri }, position: { character: 4, line: 4 } },
+        params: { textDocument: { uri: uri }, position: { character: 4, line: 6 } },
       )
-      refute_empty(server.pop_response.response)
+      assert_equal(1, server.pop_response.response.first.range.start.line)
     end
   end