Merge pull request #3878 from ruby/find-all

Add Prism::Node#find_all

Author: kddnewton

templates/lib/prism/node.rb.erb

@@ -200,7 +200,7 @@ module Prism
     end
 
     # Returns the first node that matches the given block when visited in a
-    # depth-first search. This is useful for finding a node that matches a
+    # breadth-first search. This is useful for finding a node that matches a
     # particular condition.
     #
     #     node.breadth_first_search { |node| node.node_id == node_id }
@@ -215,6 +215,26 @@ module Prism
 
       nil
     end
+    alias find breadth_first_search
+
+    # Returns all of the nodes that match the given block when visited in a
+    # breadth-first search. This is useful for finding all nodes that match a
+    # particular condition.
+    #
+    #     node.breadth_first_search_all { |node| node.is_a?(Prism::CallNode) }
+    #
+    def breadth_first_search_all(&block)
+      queue = [self] #: Array[Prism::node]
+      results = [] #: Array[Prism::node]
+
+      while (node = queue.shift)
+        results << node if yield node
+        queue.concat(node.compact_child_nodes)
+      end
+
+      results
+    end
+    alias find_all breadth_first_search_all
 
     # Returns a list of the fields that exist for this node class. Fields
     # describe the structure of the node. This kind of reflection is useful for

templates/rbi/prism/node.rbi.erb

@@ -49,6 +49,9 @@ class Prism::Node
   sig { params(block: T.proc.params(node: Prism::Node).returns(T::Boolean)).returns(T.nilable(Prism::Node)) }
   def breadth_first_search(&block); end
 
+  sig { params(block: T.proc.params(node: Prism::Node).returns(T::Boolean)).returns(T::Array[Prism::Node]) }
+  def breadth_first_search_all(&block); end
+
   sig { abstract.params(visitor: Prism::Visitor).returns(T.untyped) }
   def accept(visitor); end
 

templates/sig/prism/node.rbs.erb

@@ -25,6 +25,9 @@ module Prism
     def to_dot: () -> String
     def tunnel: (Integer line, Integer column) -> Array[Prism::node]
     def breadth_first_search: () { (Prism::node) -> bool } -> Prism::node?
+    alias find breadth_first_search
+    def breadth_first_search_all: () { (Prism::node) -> bool } -> Array[Prism::node]
+    alias find_all breadth_first_search_all
     def newline!: (Array[untyped]) -> void
 
     def save: (_Repository repository) -> void

test/prism/result/breadth_first_search_test.rb

@@ -14,5 +14,16 @@ def test_breadth_first_search
       refute_nil found
       assert_equal 8, found.start_offset
     end
+
+    def test_breadth_first_search_all
+      result = Prism.parse("[1 + 2, 2]")
+      found_nodes =
+        result.value.breadth_first_search_all do |node|
+          node.is_a?(IntegerNode)
+        end
+
+      assert_equal 3, found_nodes.size
+      assert_equal 8, found_nodes[0].start_offset
+    end
   end
 end