Merge pull request #3826 from Earlopain/optimize-ruby-visitor

Optimize ruby visitor

Author: kddnewton

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

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

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)

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

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

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) } -%>

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

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