Prism.find

lib/prism.rb

@@ -23,6 +23,7 @@ module Prism
  autoload :InspectVisitor, "prism/inspect_visitor"
  autoload :LexCompat, "prism/lex_compat"
  autoload :MutationCompiler, "prism/mutation_compiler"
+ autoload :NodeFind, "prism/node_find"
  autoload :Pattern, "prism/pattern"
  autoload :Reflection, "prism/reflection"
  autoload :Relocation, "prism/relocation"
@@ -34,7 +35,10 @@ module Prism
  # Some of these constants are not meant to be exposed, so marking them as
  # private here.
 
- private_constant :LexCompat
+ if RUBY_ENGINE != "jruby"
+ private_constant :LexCompat
+ private_constant :NodeFind
+ end
 
  # Raised when requested to parse as the currently running Ruby version but Prism has no support for it.
  class CurrentVersionError < ArgumentError
@@ -81,6 +85,16 @@ def self.load(source, serialized, freeze = false)
  Serialize.load_parse(source, serialized, freeze)
  end
 
+ # Given a Method, UnboundMethod, Proc, or Thread::Backtrace::Location,
+ # returns the Prism node representing it. On CRuby, this uses node_id for
+ # an exact match. On other implementations, it falls back to best-effort
+ # matching by source location line number.
+ #--
+ #: (Method | UnboundMethod | Proc | Thread::Backtrace::Location callable, ?rubyvm: bool) -> Node?
+ def self.find(callable, rubyvm: !!defined?(RubyVM))
+ NodeFind.find(callable, rubyvm)
+ end
+
  # @rbs!
  # VERSION: String
  # BACKEND: :CEXT | :FFI

lib/prism/node_find.rb

@@ -0,0 +1,193 @@
+# frozen_string_literal: true
+# :markup: markdown
+#--
+# rbs_inline: enabled
+
+module Prism
+ # Finds the Prism AST node corresponding to a given Method, UnboundMethod,
+ # Proc, or Thread::Backtrace::Location. On CRuby, uses node_id from the
+ # instruction sequence for an exact match. On other implementations, falls
+ # back to best-effort matching by source location line number.
+ #
+ # This module is autoloaded so that programs that don't use Prism.find don't
+ # pay for its definition.
+ module NodeFind # :nodoc:
+ # Find the node for the given callable or backtrace location.
+ #--
+ #: (Method | UnboundMethod | Proc | Thread::Backtrace::Location callable, bool rubyvm) -> Node?
+ #++
+ def self.find(callable, rubyvm)
+ case callable
+ when Proc
+ if rubyvm
+ RubyVMCallableFind.new.find(callable)
+ elsif callable.lambda?
+ LineLambdaFind.new.find(callable)
+ else
+ LineProcFind.new.find(callable)
+ end
+ when Method, UnboundMethod
+ if rubyvm
+ RubyVMCallableFind.new.find(callable)
+ else
+ LineMethodFind.new.find(callable)
+ end
+ when Thread::Backtrace::Location
+ if rubyvm
+ RubyVMBacktraceLocationFind.new.find(callable)
+ else
+ LineBacktraceLocationFind.new.find(callable)
+ end
+ else
+ raise ArgumentError, "Expected a Method, UnboundMethod, Proc, or Thread::Backtrace::Location, got #{callable.class}"
+ end
+ end
+
+ # Base class that handles parsing a file.
+ class Find
+ private
+
+ # Parse the given file path, returning a ParseResult or nil.
+ #--
+ #: (String? file) -> ParseResult?
+
+ def parse_file(file)
+ return unless file && File.readable?(file)
+ result = Prism.parse_file(file)
+ result if result.success?
+ end
+ end
+
+ # Finds the AST node for a Method, UnboundMethod, or Proc using the node_id
+ # from the instruction sequence.
+ class RubyVMCallableFind < Find
+ # Find the node for the given callable using the ISeq node_id.
+ #--
+ #: (Method | UnboundMethod | Proc callable) -> Node?
+
+ def find(callable)
+ return unless (source_location = callable.source_location)
+ return unless (result = parse_file(source_location[0]))
+ return unless (iseq = RubyVM::InstructionSequence.of(callable))
+
+ header = iseq.to_a[4]
+ return unless header[:parser] == :prism
+
+ result.value.find { |node| node.node_id == header[:node_id] }
+ end
+ end
+
+ # Finds the AST node for a Thread::Backtrace::Location using the node_id
+ # from the backtrace location.
+ class RubyVMBacktraceLocationFind < Find
+ # Find the node for the given backtrace location using node_id.
+ #--
+ #: (Thread::Backtrace::Location location) -> Node?
+
+ def find(location)
+ file = location.absolute_path || location.path
+ return unless (result = parse_file(file))
+ return unless RubyVM::AbstractSyntaxTree.respond_to?(:node_id_for_backtrace_location)
+
+ node_id = RubyVM::AbstractSyntaxTree.node_id_for_backtrace_location(location)
+
+ result.value.find { |node| node.node_id == node_id }
+ end
+ end
+
+ # Finds the AST node for a Method or UnboundMethod using best-effort line
+ # matching. Used on non-CRuby implementations.
+ class LineMethodFind < Find
+ # Find the node for the given method by matching on name and line.
+ #--
+ #: (Method | UnboundMethod callable) -> Node?
+
+ def find(callable)
+ return unless (source_location = callable.source_location)
+ return unless (result = parse_file(source_location[0]))
+
+ name = callable.name
+ start_line = source_location[1]
+
+ result.value.find do |node|
+ case node
+ when DefNode
+ node.name == name && node.location.start_line == start_line
+ when CallNode
+ node.block.is_a?(BlockNode) && node.location.start_line == start_line
+ else
+ false
+ end
+ end
+ end
+ end
+
+ # Finds the AST node for a lambda using best-effort line matching. Used
+ # on non-CRuby implementations.
+ class LineLambdaFind < Find
+ # Find the node for the given lambda by matching on line.
+ #--
+ #: (Proc callable) -> Node?
+
+ def find(callable)
+ return unless (source_location = callable.source_location)
+ return unless (result = parse_file(source_location[0]))
+
+ start_line = source_location[1]
+
+ result.value.find do |node|
+ case node
+ when LambdaNode
+ node.location.start_line == start_line
+ when CallNode
+ node.block.is_a?(BlockNode) && node.location.start_line == start_line
+ else
+ false
+ end
+ end
+ end
+ end
+
+ # Finds the AST node for a non-lambda Proc using best-effort line
+ # matching. Used on non-CRuby implementations.
+ class LineProcFind < Find
+ # Find the node for the given proc by matching on line.
+ #--
+ #: (Proc callable) -> Node?
+
+ def find(callable)
+ return unless (source_location = callable.source_location)
+ return unless (result = parse_file(source_location[0]))
+
+ start_line = source_location[1]
+
+ result.value.find do |node|
+ case node
+ when ForNode
+ node.location.start_line == start_line
+ when CallNode
+ node.block.is_a?(BlockNode) && node.location.start_line == start_line
+ else
+ false
+ end
+ end
+ end
+ end
+
+ # Finds the AST node for a Thread::Backtrace::Location using best-effort
+ # line matching. Used on non-CRuby implementations.
+ class LineBacktraceLocationFind < Find
+ # Find the node for the given backtrace location by matching on line.
+ #--
+ #: (Thread::Backtrace::Location location) -> Node?
+
+ def find(location)
+ file = location.absolute_path || location.path
+ return unless (result = parse_file(file))
+
+ start_line = location.lineno
+ result.value.find { |node| node.location.start_line == start_line }
+ end
+ end
+ end
+end

prism.gemspec

@@ -117,6 +117,7 @@ Gem::Specification.new do |spec|
  "lib/prism/lex_compat.rb",
  "lib/prism/mutation_compiler.rb",
  "lib/prism/node_ext.rb",
+ "lib/prism/node_find.rb",
  "lib/prism/node.rb",
  "lib/prism/parse_result.rb",
  "lib/prism/parse_result/comments.rb",
@@ -158,6 +159,7 @@ Gem::Specification.new do |spec|
  "rbi/generated/prism/mutation_compiler.rbi",
  "rbi/generated/prism/node.rbi",
  "rbi/generated/prism/node_ext.rbi",
+ "rbi/generated/prism/node_find.rbi",
  "rbi/generated/prism/parse_result.rbi",
  "rbi/generated/prism/pattern.rbi",
  "rbi/generated/prism/reflection.rbi",
@@ -172,6 +174,7 @@ Gem::Specification.new do |spec|
  "rbi/prism/translation/parser.rbi",
  "rbi/prism/translation/parser_versions.rbi",
  "rbi/prism/translation/ripper.rbi",
+ "rbi/rubyvm/node_find.rbi",
  "sig/generated/prism.rbs",
  "sig/generated/prism/compiler.rbs",
  "sig/generated/prism/desugar_compiler.rbs",
@@ -183,6 +186,7 @@ Gem::Specification.new do |spec|
  "sig/generated/prism/mutation_compiler.rbs",
  "sig/generated/prism/node.rbs",
  "sig/generated/prism/node_ext.rbs",
+ "sig/generated/prism/node_find.rbs",
  "sig/generated/prism/parse_result.rbs",
  "sig/generated/prism/pattern.rbs",
  "sig/generated/prism/reflection.rbs",

rakelib/typecheck.rake

@@ -358,6 +358,7 @@ namespace :typecheck do
  - ./lib/prism/lex_compat.rb
  - ./lib/prism/node.rb
  - ./lib/prism/node_ext.rb
+ - ./lib/prism/node_find.rb
  - ./lib/prism/parse_result.rb
  - ./lib/prism/pattern.rb
  - ./lib/prism/visitor.rb

rbi/rubyvm/node_find.rbi

@@ -0,0 +1,14 @@
+# typed: true
+
+class RubyVM::InstructionSequence
+ sig { params(callable: T.any(Method, UnboundMethod, Proc)).returns(T.nilable(RubyVM::InstructionSequence)) }
+ def self.of(callable); end
+
+ sig { returns(T::Array[T.untyped]) }
+ def to_a; end
+end
+
+module RubyVM::AbstractSyntaxTree
+ sig { params(location: Thread::Backtrace::Location).returns(Integer) }
+ def self.node_id_for_backtrace_location(location); end
+end

sig/_shims/rubyvm.rbs

@@ -0,0 +1,6 @@
+class RubyVM
+ class InstructionSequence
+ def self.of: (Method | UnboundMethod | Proc) -> RubyVM::InstructionSequence?
+ def to_a: () -> Array[untyped]
+ end
+end