From cc7d11c45216440b6947a012205f86897101e1fb Mon Sep 17 00:00:00 2001
From: Lukas Fittl <lukas@fittl.com>
Date: Wed, 3 Jan 2024 13:24:51 -0800
Subject: [PATCH] Treewalker: Allow passing a block with a single argument to
 walk!

This simplifies the common use case of just working with the walked node
directly, which is now sufficient more often thanks to the addition of
Node#inner= streamling modification of Node objects without going to the
parent to replace them.
---
 lib/pg_query/treewalker.rb  | 19 ++++++++++++++++---
 spec/lib/treewalker_spec.rb |  2 +-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/lib/pg_query/treewalker.rb b/lib/pg_query/treewalker.rb
index 19f0253..e62b585 100644
--- a/lib/pg_query/treewalker.rb
+++ b/lib/pg_query/treewalker.rb
@@ -1,8 +1,21 @@
 module PgQuery
   class ParserResult
-    def walk!
-      treewalker!(@tree) do |parent_node, parent_field, node, location|
-        yield(parent_node, parent_field, node, location)
+    # Walks the parse tree and calls the passed block for each contained node
+    #
+    # If you pass a block with 1 argument, you will get each node.
+    # If you pass a block with 4 arguments, you will get each parent_node, parent_field, node and location.
+    #
+    # Location uniquely identifies a given node within the parse tree. This is a stable identifier across
+    # multiple parser runs, assuming the same pg_query release and no modifications to the parse tree.
+    def walk!(&block)
+      if block.arity == 1
+        treewalker!(@tree) do |_, _, node, _|
+          yield(node)
+        end
+      else
+        treewalker!(@tree) do |parent_node, parent_field, node, location|
+          yield(parent_node, parent_field, node, location)
+        end
       end
     end
 
diff --git a/spec/lib/treewalker_spec.rb b/spec/lib/treewalker_spec.rb
index 4a174e8..3b7870c 100644
--- a/spec/lib/treewalker_spec.rb
+++ b/spec/lib/treewalker_spec.rb
@@ -36,7 +36,7 @@
 
   it 'allows recursively replacing nodes' do
     query = PgQuery.parse("SELECT * FROM tbl WHERE col::text = ANY(((ARRAY[$39, $40])::varchar[])::text[])")
-    query.walk! do |_parent_node, _parent_field, node, _location|
+    query.walk! do |node|
       next unless node.is_a?(PgQuery::Node)
       # Keep removing type casts until we hit a different class
       node.inner = node.type_cast.arg.inner while node.node == :type_cast