From 32090e2b8d4bd1c0ddfd6b4ee82927e935534eed Mon Sep 17 00:00:00 2001
From: Kevin Newton <kddnewton@gmail.com>
Date: Wed, 3 Jul 2024 08:42:28 -0400
Subject: [PATCH] [ruby/prism] Add Node#breadth_first_search

https://github.com/ruby/prism/commit/1ffb141199
---
 prism/templates/lib/prism/node.rb.erb          | 17 +++++++++++++++++
 test/prism/result/breadth_first_search_test.rb | 18 ++++++++++++++++++
 2 files changed, 35 insertions(+)
 create mode 100644 test/prism/result/breadth_first_search_test.rb

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