Use #source_range when available for Prism.find

lib/prism/node_find.rb

@@ -16,6 +16,10 @@ module NodeFind # :nodoc:
     #
     #: (Method | UnboundMethod | Proc | Thread::Backtrace::Location callable) -> Node?
     def self.find(callable)
+      if callable.respond_to?(:source_range)
+        return SourceRangeFind.new.find(callable)
+      end
+
       case callable
       when Proc
         if defined?(::RubyVM)
@@ -56,6 +60,32 @@ def parse_file(file)
       end
     end
 
+    # Finds the AST node for a Method, UnboundMethod, Proc or Thread::Backtrace::Location
+    # using the #source_range
+    class SourceRangeFind < Find
+      # Find the node for the given callable using the #source_range.
+      #
+      #: (Method | UnboundMethod | Proc | Thread::Backtrace::Location callable) -> Node?
+      def find(callable)
+        begin
+          range = callable.source_range
+        rescue ArgumentError # eval
+          return
+        end
+        return unless range
+        return unless (result = parse_file(range.absolute_path))
+
+        start_offset = result.source.byte_offset(range.start_line, range.start_column)
+        end_offset = result.source.byte_offset(range.end_line, range.end_column)
+        result.value.tunnel(range.start_line, range.start_column).reverse_each do |node|
+          if node.start_offset == start_offset && node.end_offset == end_offset
+            return node
+          end
+        end
+        nil
+      end
+    end
+
     # Finds the AST node for a Method, UnboundMethod, or Proc using the node_id
     # from the instruction sequence.
     class RubyVMCallableFind < Find

sig/_shims/source_range.rbs

@@ -0,0 +1,33 @@
+# Method#source_range, UnboundMethod#source_range, Proc#source_range and
+# Thread::Backtrace::Location#source_range are available on Ruby 4.1+
+# and are not yet part of the RBS core signatures.
+module Ruby
+  class SourceRange
+    def path: () -> String
+    def absolute_path: () -> String?
+    def start_line: () -> Integer
+    def start_column: () -> Integer
+    def end_line: () -> Integer
+    def end_column: () -> Integer
+  end
+end
+
+class Method
+  def source_range: () -> Ruby::SourceRange?
+end
+
+class UnboundMethod
+  def source_range: () -> Ruby::SourceRange?
+end
+
+class Proc
+  def source_range: () -> Ruby::SourceRange?
+end
+
+class Thread
+  class Backtrace
+    class Location
+      def source_range: () -> Ruby::SourceRange?
+    end
+  end
+end

test/prism/ruby/find_test.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 return if RUBY_ENGINE == "ruby" && RUBY_VERSION < "3.4"
-return if defined?(RubyVM::InstructionSequence) && RubyVM::InstructionSequence.compile("").to_a[4][:parser] != :prism
+return if RUBY_VERSION < "4.1" && defined?(RubyVM::InstructionSequence) && RubyVM::InstructionSequence.compile("").to_a[4][:parser] != :prism
 
 require_relative "../test_helper"
 require_relative "find_fixtures"
@@ -189,32 +189,35 @@ def test_fallback_backtrace_location
       def test_node_id_matches_iseq
         m = Fixtures::Methods.instance_method(:simple_method)
         node = Prism.find(m)
-        assert_equal node_id_of(m), node.node_id
+        assert_same_node_id m, node
       end
 
       def test_node_id_for_lambda
         node = Prism.find(Fixtures::Procs::SIMPLE_LAMBDA)
-        assert_equal node_id_of(Fixtures::Procs::SIMPLE_LAMBDA), node.node_id
+        assert_same_node_id Fixtures::Procs::SIMPLE_LAMBDA, node
       end
 
       def test_node_id_for_proc
         node = Prism.find(Fixtures::Procs::SIMPLE_PROC)
-        assert_equal node_id_of(Fixtures::Procs::SIMPLE_PROC), node.node_id
+        assert_same_node_id Fixtures::Procs::SIMPLE_PROC, node
       end
 
       def test_node_id_for_define_method
         m = Fixtures::DefineMethod.instance_method(:dynamic)
         node = Prism.find(m)
-        assert_equal node_id_of(m), node.node_id
+        assert_same_node_id m, node
       end
 
       def test_node_id_for_backtrace_location
         location = zero_division_location
         assert_not_nil location
-        expected_node_id = RubyVM::AbstractSyntaxTree.node_id_for_backtrace_location(location)
 
-        node = Prism.find(location)
-        assert_equal expected_node_id, node.node_id
+        if RubyVM::InstructionSequence.compile("").to_a[4][:parser] == :prism
+          expected_node_id = RubyVM::AbstractSyntaxTree.node_id_for_backtrace_location(location)
+
+          node = Prism.find(location)
+          assert_equal expected_node_id, node.node_id
+        end
       end
     end
 
@@ -235,8 +238,11 @@ def zero_division_location
       fixture_backtrace_location(e)
     end
 
-    def node_id_of(callable)
-      RubyVM::InstructionSequence.of(callable).to_a[4][:node_id]
+    def assert_same_node_id(callable, node)
+      data = RubyVM::InstructionSequence.of(callable).to_a[4]
+      if data[:parser] == :prism
+        assert_equal data[:node_id], node.node_id
+      end
     end
   end
 end