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

@@ -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
 

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