diff options
| -rw-r--r-- | prism/templates/lib/prism/node.rb.erb | 17 | ||||
| -rw-r--r-- | test/prism/result/breadth_first_search_test.rb | 18 |
2 files changed, 35 insertions, 0 deletions
diff --git a/prism/templates/lib/prism/node.rb.erb b/prism/templates/lib/prism/node.rb.erb index ce166dda65..66cbce767e 100644 --- a/prism/templates/lib/prism/node.rb.erb +++ b/prism/templates/lib/prism/node.rb.erb @@ -116,6 +116,23 @@ module Prism result 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 + # particular condition. + # + # node.breadth_first_search { |node| node.node_id == node_id } + # + def breadth_first_search(&block) + queue = [self] #: Array[Prism::node] + + while (node = queue.shift) + return node if yield node + queue.concat(node.compact_child_nodes) + end + + nil + end + # 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 # things like recursively visiting each node _and_ field in the tree. diff --git a/test/prism/result/breadth_first_search_test.rb b/test/prism/result/breadth_first_search_test.rb new file mode 100644 index 0000000000..e2e043a902 --- /dev/null +++ b/test/prism/result/breadth_first_search_test.rb @@ -0,0 +1,18 @@ +# frozen_string_literal: true + +require_relative "../test_helper" + +module Prism + class BreadthFirstSearchTest < TestCase + def test_breadth_first_search + result = Prism.parse("[1 + 2, 2]") + found = + result.value.breadth_first_search do |node| + node.is_a?(IntegerNode) && node.value == 2 + end + + refute_nil found + assert_equal 8, found.start_offset + end + end +end |