From 036e668097a1ff80bcd52dce1f5585ded7924afd Mon Sep 17 00:00:00 2001
From: Lukas Fittl <lukas@fittl.com>
Date: Fri, 26 Jan 2024 22:19:58 -0800
Subject: [PATCH] Treewalker: Improve performance by avoiding memory
 allocations

Previously the treewalker would unnecessary call "to_h" on each Protobuf
message in the parsetree, in order to get the field names to walk. This
caused unnecessary copies of the message, increasing memory usage and
slowing down the tree walk. Instead, use the Protobuf descriptor and its
field descriptors to walk the message.

Additionally this also optimizes the case where a block with 1 argument
is used for the tree walk, since we don't need to handle the location,
avoiding unnecessary copies of the field name string.

Together these changes result in about a 5x speed up in some use cases.
---
 lib/pg_query/treewalker.rb | 43 ++++++++++++++++++++++++++++++++------
 1 file changed, 37 insertions(+), 6 deletions(-)

diff --git a/lib/pg_query/treewalker.rb b/lib/pg_query/treewalker.rb
index e62b585..386c310 100644
--- a/lib/pg_query/treewalker.rb
+++ b/lib/pg_query/treewalker.rb
@@ -5,15 +5,17 @@ class ParserResult
     # 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.
     #
+    # If sufficent for the use case, the 1 argument block approach is recommended, since its faster.
+    #
     # 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, _|
+        treewalker!(@tree) do |node|
           yield(node)
         end
       else
-        treewalker!(@tree) do |parent_node, parent_field, node, location|
+        treewalker_with_location!(@tree) do |parent_node, parent_field, node, location|
           yield(parent_node, parent_field, node, location)
         end
       end
@@ -22,6 +24,34 @@ def walk!(&block)
     private
 
     def treewalker!(tree) # rubocop:disable Metrics/CyclomaticComplexity
+      nodes = [tree.dup]
+
+      loop do
+        parent_node = nodes.shift
+
+        case parent_node
+        when Google::Protobuf::MessageExts
+          parent_node.class.descriptor.each do |field_descriptor|
+            node = field_descriptor.get(parent_node)
+            next if node.nil?
+            yield(node) if node.is_a?(Google::Protobuf::MessageExts) || node.is_a?(Google::Protobuf::RepeatedField)
+
+            nodes << node unless node.nil?
+          end
+        when Google::Protobuf::RepeatedField
+          parent_node.each do |node|
+            next if node.nil?
+            yield(node) if node.is_a?(Google::Protobuf::MessageExts) || node.is_a?(Google::Protobuf::RepeatedField)
+
+            nodes << node unless node.nil?
+          end
+        end
+
+        break if nodes.empty?
+      end
+    end
+
+    def treewalker_with_location!(tree) # rubocop:disable Metrics/CyclomaticComplexity
       nodes = [[tree.dup, []]]
 
       loop do
@@ -29,11 +59,12 @@ def treewalker!(tree) # rubocop:disable Metrics/CyclomaticComplexity
 
         case parent_node
         when Google::Protobuf::MessageExts
-          parent_node.to_h.keys.each do |parent_field|
-            node = parent_node[parent_field.to_s]
+          parent_node.class.descriptor.each do |field_descriptor|
+            parent_field = field_descriptor.name
+            node = parent_node[parent_field]
             next if node.nil?
-            location = parent_location + [parent_field]
-            yield(parent_node, parent_field, node, location) if node.is_a?(Google::Protobuf::MessageExts) || node.is_a?(Google::Protobuf::RepeatedField)
+            location = parent_location + [parent_field.to_sym]
+            yield(parent_node, parent_field.to_sym, node, location) if node.is_a?(Google::Protobuf::MessageExts) || node.is_a?(Google::Protobuf::RepeatedField)
 
             nodes << [node, location] unless node.nil?
           end