Add Prism.node_for(Method|UnboundMethod|Proc) to get a Prism node for a callable

* See https://bugs.ruby-lang.org/issues/21005
  and https://bugs.ruby-lang.org/issues/20999

Author: eregon

lib/prism/parse_result.rb

@@ -143,6 +143,13 @@ def code_units_column(byte_offset, encoding)
       code_units_offset(byte_offset, encoding) - code_units_offset(line_start(byte_offset), encoding)
     end
 
+    # Returns the byte offset for a given line number and column number
+    def line_and_column_to_byte_offset(line, column)
+      line_start = offsets[line - 1]
+      line_end = offsets[line]
+      line_start + @source.byteslice(line_start, line_end)[0...column].bytesize
+    end
+
     # Freeze this object and the objects it contains.
     def deep_freeze
       source.freeze
@@ -272,6 +279,13 @@ def code_units_cache(encoding)
     def code_units_column(byte_offset, encoding)
       byte_offset - line_start(byte_offset)
     end
+
+    # Specialized version of `line_and_column_to_byte_offset`
+    # which does not need to access the source String
+    def line_and_column_to_byte_offset(line, column)
+      line_start = offsets[line - 1]
+      line_start + column
+    end
   end
 
   # This represents a location in the source.
@@ -895,4 +909,46 @@ def initialize(locals, forwarding)
   def self.scope(locals: [], forwarding: [])
     Scope.new(locals, forwarding)
   end
+
+  # Given a Method, UnboundMethod or Proc, use its #source_location to parse the file and return a Prism node representing it.
+  # The returned node will either be a DefNode, LambdaNode or CallNode.
+  # Raises ArgumentError if it cannot be found for any reason.
+  # Only works on Ruby 4+ as it needs #source_location to contain column and end line information.
+  def self.node_for(callable)
+    unless callable.is_a?(Method) || callable.is_a?(UnboundMethod) || callable.is_a?(Proc)
+      raise ArgumentError, 'Prism.node_for requires a Method, UnboundMethod or Proc'
+    end
+    source_location = callable.source_location
+    raise ArgumentError, "#source_location is nil for #{callable}" if source_location.nil?
+    raise ArgumentError, '#source_location does not contain column and end_line, this method only works on Ruby 4+' if source_location.size != 5
+    file, start_line, start_column, end_line, end_column = source_location
+
+    unless File.exist?(file)
+      raise ArgumentError, "#source_location[0] is #{file} but this file does not exist"
+    end
+
+    parse_result = Prism.parse_file(file)
+    unless parse_result.success?
+      raise ArgumentError, "#{file} has syntax errors: #{parse_result.errors_format}"
+    end
+    root = parse_result.value
+    start_offset = parse_result.source.line_and_column_to_byte_offset(start_line, start_column)
+    end_offset = parse_result.source.line_and_column_to_byte_offset(end_line, end_column)
+
+    found = root.breadth_first_search do |node|
+      case node
+      when DefNode
+        node.start_offset == start_offset && node.end_offset == end_offset
+      when LambdaNode
+        # Proc#source_location returns start_column 2 for `-> { ... }` (just after the `->`)
+        node.operator_loc.end_offset == start_offset && node.end_offset == end_offset
+      when CallNode
+        # Proc#source_location returns start_column 5 for `proc { ... }` (the `{`)
+        node.block && node.block.opening_loc.start_offset == start_offset && node.end_offset == end_offset
+      end
+    end
+
+    raise ArgumentError, "Could not find node for #{callable} in #{file} at (#{start_line},#{start_column})-(#{end_line},#{end_column})" unless found
+    found
+  end
 end

sig/prism/parse_result.rbs

@@ -24,6 +24,7 @@ module Prism
     def code_units_offset: (Integer byte_offset, Encoding encoding) -> Integer
     def code_units_cache: (Encoding encoding) -> _CodeUnitsCache
     def code_units_column: (Integer byte_offset, Encoding encoding) -> Integer
+    def line_and_column_to_byte_offset: (Integer line, Integer column) -> Integer
     def deep_freeze: () -> void
 
     def self.for: (String source) -> Source
@@ -40,6 +41,7 @@ module Prism
     def code_units_offset: (Integer byte_offset, Encoding encoding) -> Integer
     def code_units_cache: (Encoding encoding) -> _CodeUnitsCache
     def code_units_column: (Integer byte_offset, Encoding encoding) -> Integer
+    def line_and_column_to_byte_offset: (Integer line, Integer column) -> Integer
   end
 
   class Location

templates/sig/prism.rbs.erb

@@ -90,4 +90,6 @@ module Prism
   ) -> ParseResult
 
   def self.scope: (?locals: Array[Symbol], ?forwarding: Array[Symbol]) -> Scope
+
+  def self.node_for: (callable: Method | UnboundMethod | Proc) -> (DefNode | LambdaNode | CallNode | nil)
 end

test/prism/ruby/node_for_test.rb

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+# typed: ignore
+
+require_relative "../test_helper"
+
+module Prism
+  class NodeForTest < TestCase
+    INDENT = ' ' * 4
+
+    def m(foo)
+      42
+    end
+
+    def inline_method = 42
+
+    define_method(:define_method_method) { 42 }
+
+    def return_block(&block) = block
+
+    def test_def_method
+      node = Prism.node_for(NodeForTest.instance_method(:m))
+      assert_instance_of(Prism::DefNode, node)
+      assert_equal "def m(foo)\n#{INDENT}  42\n#{INDENT}end", node.slice
+
+      node = Prism.node_for(method(:m))
+      assert_instance_of(Prism::DefNode, node)
+      assert_equal "def m(foo)\n      42\n    end", node.slice
+    end
+
+    def test_inline_method
+      node = Prism.node_for(method(:inline_method))
+      assert_instance_of(Prism::DefNode, node)
+      assert_equal "def inline_method = 42", node.slice
+    end
+
+    def test_define_method
+      node = Prism.node_for(method(:define_method_method))
+      assert_instance_of(Prism::CallNode, node)
+      assert_equal "define_method(:define_method_method) { 42 }", node.slice
+      assert_equal "{ 42 }", node.block.slice
+    end
+
+    def test_lambda
+      node = Prism.node_for(-> { 42 })
+      assert_instance_of(Prism::LambdaNode, node)
+      assert_equal "-> { 42 }", node.slice
+      assert_equal "{ 42 }", node.opening_loc.join(node.closing_loc).slice
+
+      node = Prism.node_for(lambda { 42 })
+      assert_instance_of(Prism::CallNode, node)
+      assert_equal "lambda { 42 }", node.slice
+      assert_equal "{ 42 }", node.block.slice
+    end
+
+    def test_proc
+      node = Prism.node_for(proc { 42 })
+      assert_instance_of(Prism::CallNode, node)
+      assert_equal "proc { 42 }", node.slice
+      assert_equal "{ 42 }", node.block.slice
+
+      node = Prism.node_for(return_block { 42 })
+      assert_instance_of(Prism::CallNode, node)
+      assert_equal "return_block { 42 }", node.slice
+      assert_equal "{ 42 }", node.block.slice
+
+      heredoc_proc = proc { <<~END }
+        heredoc
+      END
+      node = Prism.node_for(heredoc_proc)
+      assert_instance_of(Prism::CallNode, node)
+      assert_equal "proc { <<~END }", node.slice
+      assert_equal "heredoc\n", node.block.body.body.first.unescaped
+    end
+
+    def test_method_to_proc
+      node = Prism.node_for(method(:inline_method).to_proc)
+      assert_instance_of(Prism::DefNode, node)
+      assert_equal "def inline_method = 42", node.slice
+    end
+
+    def test_eval
+      l = eval "-> { 42 }"
+      e = assert_raise(ArgumentError) { Prism.node_for(l) }
+      assert_include e.message, 'eval'
+
+      l = eval "-> { 42 }", nil, __FILE__, __LINE__
+      e = assert_raise(ArgumentError) { Prism.node_for(l) }
+      assert_include e.message, 'Could not find node'
+    end
+  end
+end