Skip to content

Commit 62280d2

Browse files
committed
Prism.find
Take a method, unbound method, proc, or thread backtrace location. This is our equivalent to RubyVM::AbstractSyntaxTree.of, and could be leveraged in something like error highlight. Note that this uses CRuby-specific APIs on CRuby, and falls back to using location-based APIs when those aren't available.
1 parent f95bef9 commit 62280d2

11 files changed

Lines changed: 683 additions & 1 deletion

File tree

lib/prism.rb

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ module Prism
2323
autoload :InspectVisitor, "prism/inspect_visitor"
2424
autoload :LexCompat, "prism/lex_compat"
2525
autoload :MutationCompiler, "prism/mutation_compiler"
26+
autoload :NodeFind, "prism/node_find"
2627
autoload :Pattern, "prism/pattern"
2728
autoload :Reflection, "prism/reflection"
2829
autoload :Relocation, "prism/relocation"
@@ -34,7 +35,10 @@ module Prism
3435
# Some of these constants are not meant to be exposed, so marking them as
3536
# private here.
3637

37-
private_constant :LexCompat
38+
if RUBY_ENGINE != "jruby"
39+
private_constant :LexCompat
40+
private_constant :NodeFind
41+
end
3842

3943
# Raised when requested to parse as the currently running Ruby version but Prism has no support for it.
4044
class CurrentVersionError < ArgumentError
@@ -81,6 +85,16 @@ def self.load(source, serialized, freeze = false)
8185
Serialize.load_parse(source, serialized, freeze)
8286
end
8387

88+
# Given a Method, UnboundMethod, Proc, or Thread::Backtrace::Location,
89+
# returns the Prism node representing it. On CRuby, this uses node_id for
90+
# an exact match. On other implementations, it falls back to best-effort
91+
# matching by source location line number.
92+
#--
93+
#: (Method | UnboundMethod | Proc | Thread::Backtrace::Location callable, ?rubyvm: bool) -> Node?
94+
def self.find(callable, rubyvm: defined?(RubyVM))
95+
NodeFind.find(callable, rubyvm)
96+
end
97+
8498
# @rbs!
8599
# VERSION: String
86100
# BACKEND: :CEXT | :FFI

lib/prism/node_find.rb

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
# frozen_string_literal: true
2+
# :markup: markdown
3+
#--
4+
# rbs_inline: enabled
5+
6+
module Prism
7+
# Finds the Prism AST node corresponding to a given Method, UnboundMethod,
8+
# Proc, or Thread::Backtrace::Location. On CRuby, uses node_id from the
9+
# instruction sequence for an exact match. On other implementations, falls
10+
# back to best-effort matching by source location line number.
11+
#
12+
# This module is autoloaded so that programs that don't use Prism.find don't
13+
# pay for its definition.
14+
module NodeFind # :nodoc:
15+
# Find the node for the given callable or backtrace location.
16+
#--
17+
#: (Method | UnboundMethod | Proc | Thread::Backtrace::Location callable, bool rubyvm) -> Node?
18+
#++
19+
def self.find(callable, rubyvm)
20+
case callable
21+
when Proc
22+
if rubyvm
23+
RubyVMCallableFind.new.find(callable)
24+
elsif callable.lambda?
25+
LineLambdaFind.new.find(callable)
26+
else
27+
LineProcFind.new.find(callable)
28+
end
29+
when Method, UnboundMethod
30+
if rubyvm
31+
RubyVMCallableFind.new.find(callable)
32+
else
33+
LineMethodFind.new.find(callable)
34+
end
35+
when Thread::Backtrace::Location
36+
if rubyvm
37+
RubyVMBacktraceLocationFind.new.find(callable)
38+
else
39+
LineBacktraceLocationFind.new.find(callable)
40+
end
41+
else
42+
raise ArgumentError, "Expected a Method, UnboundMethod, Proc, or Thread::Backtrace::Location, got #{callable.class}"
43+
end
44+
end
45+
46+
# Base class that handles parsing a file.
47+
class Find
48+
private
49+
50+
# Parse the given file path, returning a ParseResult or nil.
51+
#--
52+
#: (String? file) -> ParseResult?
53+
54+
def parse_file(file)
55+
return unless file && File.readable?(file)
56+
result = Prism.parse_file(file)
57+
result if result.success?
58+
end
59+
end
60+
61+
# Finds the AST node for a Method, UnboundMethod, or Proc using the node_id
62+
# from the instruction sequence.
63+
class RubyVMCallableFind < Find
64+
# Find the node for the given callable using the ISeq node_id.
65+
#--
66+
#: (Method | UnboundMethod | Proc callable) -> Node?
67+
68+
def find(callable)
69+
return unless (source_location = callable.source_location)
70+
return unless (result = parse_file(source_location[0]))
71+
return unless (iseq = RubyVM::InstructionSequence.of(callable))
72+
73+
header = iseq.to_a[4]
74+
return unless header[:parser] == :prism
75+
76+
result.value.find { |node| node.node_id == header[:node_id] }
77+
end
78+
end
79+
80+
# Finds the AST node for a Thread::Backtrace::Location using the node_id
81+
# from the backtrace location.
82+
class RubyVMBacktraceLocationFind < Find
83+
# Find the node for the given backtrace location using node_id.
84+
#--
85+
#: (Thread::Backtrace::Location location) -> Node?
86+
87+
def find(location)
88+
file = location.absolute_path || location.path
89+
return unless (result = parse_file(file))
90+
return unless RubyVM::AbstractSyntaxTree.respond_to?(:node_id_for_backtrace_location)
91+
92+
node_id = RubyVM::AbstractSyntaxTree.node_id_for_backtrace_location(location)
93+
94+
result.value.find { |node| node.node_id == node_id }
95+
end
96+
end
97+
98+
# Finds the AST node for a Method or UnboundMethod using best-effort line
99+
# matching. Used on non-CRuby implementations.
100+
class LineMethodFind < Find
101+
# Find the node for the given method by matching on name and line.
102+
#--
103+
#: (Method | UnboundMethod callable) -> Node?
104+
105+
def find(callable)
106+
return unless (source_location = callable.source_location)
107+
return unless (result = parse_file(source_location[0]))
108+
109+
name = callable.name
110+
start_line = source_location[1]
111+
112+
result.value.find do |node|
113+
case node
114+
when DefNode
115+
node.name == name && node.location.start_line == start_line
116+
when CallNode
117+
node.block.is_a?(BlockNode) && node.location.start_line == start_line
118+
else
119+
false
120+
end
121+
end
122+
end
123+
end
124+
125+
# Finds the AST node for a lambda using best-effort line matching. Used
126+
# on non-CRuby implementations.
127+
class LineLambdaFind < Find
128+
# Find the node for the given lambda by matching on line.
129+
#--
130+
#: (Proc callable) -> Node?
131+
132+
def find(callable)
133+
return unless (source_location = callable.source_location)
134+
return unless (result = parse_file(source_location[0]))
135+
136+
start_line = source_location[1]
137+
138+
result.value.find do |node|
139+
case node
140+
when LambdaNode
141+
node.location.start_line == start_line
142+
when CallNode
143+
node.block.is_a?(BlockNode) && node.location.start_line == start_line
144+
else
145+
false
146+
end
147+
end
148+
end
149+
end
150+
151+
# Finds the AST node for a non-lambda Proc using best-effort line
152+
# matching. Used on non-CRuby implementations.
153+
class LineProcFind < Find
154+
# Find the node for the given proc by matching on line.
155+
#--
156+
#: (Proc callable) -> Node?
157+
158+
def find(callable)
159+
return unless (source_location = callable.source_location)
160+
return unless (result = parse_file(source_location[0]))
161+
162+
start_line = source_location[1]
163+
164+
result.value.find do |node|
165+
case node
166+
when ForNode
167+
node.location.start_line == start_line
168+
when CallNode
169+
node.block.is_a?(BlockNode) && node.location.start_line == start_line
170+
else
171+
false
172+
end
173+
end
174+
end
175+
end
176+
177+
# Finds the AST node for a Thread::Backtrace::Location using best-effort
178+
# line matching. Used on non-CRuby implementations.
179+
class LineBacktraceLocationFind < Find
180+
# Find the node for the given backtrace location by matching on line.
181+
#--
182+
#: (Thread::Backtrace::Location location) -> Node?
183+
184+
def find(location)
185+
file = location.absolute_path || location.path
186+
return unless (result = parse_file(file))
187+
188+
start_line = location.lineno
189+
result.value.find { |node| node.location.start_line == start_line }
190+
end
191+
end
192+
end
193+
end

prism.gemspec

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ Gem::Specification.new do |spec|
117117
"lib/prism/lex_compat.rb",
118118
"lib/prism/mutation_compiler.rb",
119119
"lib/prism/node_ext.rb",
120+
"lib/prism/node_find.rb",
120121
"lib/prism/node.rb",
121122
"lib/prism/parse_result.rb",
122123
"lib/prism/parse_result/comments.rb",
@@ -158,6 +159,7 @@ Gem::Specification.new do |spec|
158159
"rbi/generated/prism/mutation_compiler.rbi",
159160
"rbi/generated/prism/node.rbi",
160161
"rbi/generated/prism/node_ext.rbi",
162+
"rbi/generated/prism/node_find.rbi",
161163
"rbi/generated/prism/parse_result.rbi",
162164
"rbi/generated/prism/pattern.rbi",
163165
"rbi/generated/prism/reflection.rbi",
@@ -183,6 +185,7 @@ Gem::Specification.new do |spec|
183185
"sig/generated/prism/mutation_compiler.rbs",
184186
"sig/generated/prism/node.rbs",
185187
"sig/generated/prism/node_ext.rbs",
188+
"sig/generated/prism/node_find.rbs",
186189
"sig/generated/prism/parse_result.rbs",
187190
"sig/generated/prism/pattern.rbs",
188191
"sig/generated/prism/reflection.rbs",

rbi/generated/prism.rbi

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

rbi/generated/prism/node_find.rbi

Lines changed: 64 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

sig/_shims/rubyvm.rbs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
class RubyVM
2+
class InstructionSequence
3+
def self.of: (Method | UnboundMethod | Proc) -> RubyVM::InstructionSequence?
4+
def to_a: () -> Array[untyped]
5+
end
6+
end

sig/generated/prism.rbs

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)