From bf631750cf39827b750f13d55c90af0ea95817ea Mon Sep 17 00:00:00 2001 From: Earlopain <14981592+Earlopain@users.noreply.github.com> Date: Tue, 23 Dec 2025 14:41:09 +0100 Subject: [PATCH] Optimize ruby visitor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `compact_child_nodes` allocates an array. We can skip that step by simply yielding the nodes. Benchmark for visiting the rails codebase: ```rb require "prism" require "benchmark/ips" files = Dir.glob("../rails/**/*.rb") results = files.map { Prism.parse_file(it) } visitor = Prism::Visitor.new Benchmark.ips do |x| x.config(warmup: 3, time: 10) x.report do results.each do visitor.visit(it.value) end end end RubyVM::YJIT.enable Benchmark.ips do |x| x.config(warmup: 3, time: 10) x.report do results.each do visitor.visit(it.value) end end end ``` Before: ``` ruby 3.4.8 (2025-12-17 revision 995b59f666) +PRISM [x86_64-linux] Warming up -------------------------------------- 1.000 i/100ms Calculating ------------------------------------- 2.691 (± 0.0%) i/s (371.55 ms/i) - 27.000 in 10.089422s ruby 3.4.8 (2025-12-17 revision 995b59f666) +YJIT +PRISM [x86_64-linux] Warming up -------------------------------------- 1.000 i/100ms Calculating ------------------------------------- 7.278 (±13.7%) i/s (137.39 ms/i) - 70.000 in 10.071568s ``` After: ``` ruby 3.4.8 (2025-12-17 revision 995b59f666) +PRISM [x86_64-linux] Warming up -------------------------------------- 1.000 i/100ms Calculating ------------------------------------- 3.429 (± 0.0%) i/s (291.65 ms/i) - 35.000 in 10.208580s ruby 3.4.8 (2025-12-17 revision 995b59f666) +YJIT +PRISM [x86_64-linux] Warming up -------------------------------------- 1.000 i/100ms Calculating ------------------------------------- 16.815 (± 0.0%) i/s (59.47 ms/i) - 169.000 in 10.054668s ``` ~21% faster on the interpreter, ~56% with YJIT --- docs/ruby_api.md | 1 + lib/prism/translation/ruby_parser.rb | 2 +- sample/prism/find_comments.rb | 4 ++-- sample/prism/locate_nodes.rb | 11 ++++++----- templates/lib/prism/compiler.rb.erb | 4 ++-- templates/lib/prism/node.rb.erb | 25 ++++++++++++++++++++++++- templates/lib/prism/visitor.rb.erb | 4 ++-- templates/sig/prism/node.rbs.erb | 1 + 8 files changed, 39 insertions(+), 13 deletions(-) diff --git a/docs/ruby_api.md b/docs/ruby_api.md index 6cfde4c1f3..c378417cfc 100644 --- a/docs/ruby_api.md +++ b/docs/ruby_api.md @@ -36,6 +36,7 @@ Once you have nodes in hand coming out of a parse result, there are a number of * `#accept(visitor)` - a method that will immediately call `visit_*` to specialize for the node type * `#child_nodes` - a positional array of the child nodes of the node, with `nil` values for any missing children * `#compact_child_nodes` - a positional array of the child nodes of the node with no `nil` values +* `#each_child_node` - with a block given yields all child nodes, without a block return an enumerator containing all child nodes * `#copy(**keys)` - a method that allows creating a shallow copy of the node with the given keys overridden * `#deconstruct`/`#deconstruct_keys(keys)` - the pattern matching interface for nodes * `#inspect` - a string representation that looks like the syntax tree of the node diff --git a/lib/prism/translation/ruby_parser.rb b/lib/prism/translation/ruby_parser.rb index 1fb0e27846..c026c4ad9c 100644 --- a/lib/prism/translation/ruby_parser.rb +++ b/lib/prism/translation/ruby_parser.rb @@ -1422,7 +1422,7 @@ def visit_or_node(node) # ``` def visit_parameters_node(node) children = - node.compact_child_nodes.map do |element| + node.each_child_node.map do |element| if element.is_a?(MultiTargetNode) visit_destructured_parameter(element) else diff --git a/sample/prism/find_comments.rb b/sample/prism/find_comments.rb index 6a26cd32b7..314bbb9b8e 100644 --- a/sample/prism/find_comments.rb +++ b/sample/prism/find_comments.rb @@ -16,7 +16,7 @@ def initialize(target, comments, nesting = []) # method definition on the node. def visit_module_node(node) visitor = FindMethodComments.new(@target, @comments, [*@nesting, node.name]) - node.compact_child_nodes.each { |child| child.accept(visitor) } + node.each_child_node { |child| child.accept(visitor) } end def visit_class_node(node) @@ -26,7 +26,7 @@ def visit_class_node(node) # because then the state is immutable and it's easier to reason about. This # also provides for more debugging opportunity in the initializer. visitor = FindMethodComments.new(@target, @comments, [*@nesting, node.name]) - node.compact_child_nodes.each { |child| child.accept(visitor) } + node.each_child_node { |child| child.accept(visitor) } end def visit_def_node(node) diff --git a/sample/prism/locate_nodes.rb b/sample/prism/locate_nodes.rb index 7a51db4367..dd70f87850 100644 --- a/sample/prism/locate_nodes.rb +++ b/sample/prism/locate_nodes.rb @@ -40,11 +40,12 @@ def locate(node, line:, column:) while (node = queue.shift) result << node - # Nodes have `child_nodes` and `compact_child_nodes`. `child_nodes` have - # consistent indices but include `nil` for optional fields that are not - # present, whereas `compact_child_nodes` has inconsistent indices but does - # not include `nil` for optional fields that are not present. - node.compact_child_nodes.find do |child| + # Nodes have `child_nodes`, `compact_child_nodes`, and `each_child_node`. + # `child_nodes` have consistent indices but include `nil` for optional fields that + # are not present, whereas `compact_child_nodes` has inconsistent indices but does + # not include `nil` for optional fields that are not present. `each_child_node` is + # like `compact_child_nodes`, returning an enumerator instead of an array. + node.each_child_node.find do |child| queue << child if covers?(child.location, line: line, column: column) end end diff --git a/templates/lib/prism/compiler.rb.erb b/templates/lib/prism/compiler.rb.erb index 9102025c20..66dbe666b9 100644 --- a/templates/lib/prism/compiler.rb.erb +++ b/templates/lib/prism/compiler.rb.erb @@ -29,14 +29,14 @@ module Prism # Visit the child nodes of the given node. def visit_child_nodes(node) - node.compact_child_nodes.map { |node| node.accept(self) } + node.each_child_node.map { |node| node.accept(self) } end <%- nodes.each_with_index do |node, index| -%> <%= "\n" if index != 0 -%> # Compile a <%= node.name %> node def visit_<%= node.human %>(node) - node.compact_child_nodes.map { |node| node.accept(self) } + node.each_child_node.map { |node| node.accept(self) } end <%- end -%> end diff --git a/templates/lib/prism/node.rb.erb b/templates/lib/prism/node.rb.erb index ceee2b0ffe..c97c029d3b 100644 --- a/templates/lib/prism/node.rb.erb +++ b/templates/lib/prism/node.rb.erb @@ -187,7 +187,7 @@ module Prism while (node = queue.shift) result << node - node.compact_child_nodes.each do |child_node| + node.each_child_node do |child_node| child_location = child_node.location start_line = child_location.start_line @@ -259,6 +259,13 @@ module Prism alias deconstruct child_nodes + # With a block given, yields each child node. Without a block, returns + # an enumerator that contains each child node. Excludes any `nil`s in + # the place of optional nodes that were not present. + def each_child_node + raise NoMethodError, "undefined method `each_child_node' for #{inspect}" + end + # Returns an array of child nodes, excluding any `nil`s in the place of # optional nodes that were not present. def compact_child_nodes @@ -335,6 +342,22 @@ module Prism }.compact.join(", ") %>] end + # def each_child_node: () { (Prism::node) -> void } -> void | () -> Enumerator[Prism::node] + def each_child_node + return to_enum(:each_child_node) unless block_given? + + <%- node.fields.each do |field| -%> + <%- case field -%> + <%- when Prism::Template::NodeField -%> + yield <%= field.name %> + <%- when Prism::Template::OptionalNodeField -%> + yield <%= field.name %> if <%= field.name %> + <%- when Prism::Template::NodeListField -%> + <%= field.name %>.each {|node| yield node } + <%- end -%> + <%- end -%> + end + # def compact_child_nodes: () -> Array[Node] def compact_child_nodes <%- if node.fields.any? { |field| field.is_a?(Prism::Template::OptionalNodeField) } -%> diff --git a/templates/lib/prism/visitor.rb.erb b/templates/lib/prism/visitor.rb.erb index b1a03c3f1a..76f907724f 100644 --- a/templates/lib/prism/visitor.rb.erb +++ b/templates/lib/prism/visitor.rb.erb @@ -20,7 +20,7 @@ module Prism # Visits the child nodes of `node` by calling `accept` on each one. def visit_child_nodes(node) # @type self: _Visitor - node.compact_child_nodes.each { |node| node.accept(self) } + node.each_child_node { |node| node.accept(self) } end end @@ -48,7 +48,7 @@ module Prism <%= "\n" if index != 0 -%> # Visit a <%= node.name %> node def visit_<%= node.human %>(node) - node.compact_child_nodes.each { |node| node.accept(self) } + node.each_child_node { |node| node.accept(self) } end <%- end -%> end diff --git a/templates/sig/prism/node.rbs.erb b/templates/sig/prism/node.rbs.erb index 6d6fe67336..01914b571c 100644 --- a/templates/sig/prism/node.rbs.erb +++ b/templates/sig/prism/node.rbs.erb @@ -12,6 +12,7 @@ module Prism def child_nodes: () -> Array[Prism::node?] def comment_targets: () -> Array[Prism::node | Location] def compact_child_nodes: () -> Array[Prism::node] + def each_child_node: () { (Prism::node) -> void } -> void | () -> Enumerator[Prism::node] def self.fields: () -> Array[Prism::Reflection::Field] def type: () -> Symbol def self.type: () -> Symbol