From f23f1dfe39c67d8513f3e7ebfacfeb4d9235327e Mon Sep 17 00:00:00 2001 From: Lukas Fittl Date: Mon, 29 Jan 2024 19:16:36 +0000 Subject: [PATCH] Treewalker: Improve performance by avoiding memory allocations (#309) 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/param_refs.rb | 2 +- lib/pg_query/treewalker.rb | 49 ++++++++++++++++++++++++++++--------- lib/pg_query/truncate.rb | 2 +- spec/lib/treewalker_spec.rb | 2 +- 4 files changed, 41 insertions(+), 14 deletions(-) diff --git a/lib/pg_query/param_refs.rb b/lib/pg_query/param_refs.rb index ab393f5d..0a9bd74a 100644 --- a/lib/pg_query/param_refs.rb +++ b/lib/pg_query/param_refs.rb @@ -3,7 +3,7 @@ class ParserResult def param_refs # rubocop:disable Metrics/CyclomaticComplexity results = [] - treewalker! @tree do |_, _, node, location| + treewalker_with_location! @tree do |_, _, node, location| case node when PgQuery::ParamRef # Ignore param refs inside type casts, as these are already handled diff --git a/lib/pg_query/treewalker.rb b/lib/pg_query/treewalker.rb index e62b5858..4ee9d3c1 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 sufficient for the use case, the 1 argument block approach is recommended, since it's 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,32 @@ 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 + 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 + end + end + + break if nodes.empty? + end + end + + def treewalker_with_location!(tree) # rubocop:disable Metrics/CyclomaticComplexity nodes = [[tree.dup, []]] loop do @@ -29,21 +57,20 @@ 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) - - nodes << [node, location] unless node.nil? + 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] end when Google::Protobuf::RepeatedField parent_node.each_with_index do |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) - - nodes << [node, location] unless node.nil? + nodes << [node, location] end end @@ -52,7 +79,7 @@ def treewalker!(tree) # rubocop:disable Metrics/CyclomaticComplexity end def find_tree_location(tree, searched_location) - treewalker! tree do |parent_node, parent_field, node, location| + treewalker_with_location! tree do |parent_node, parent_field, node, location| next unless location == searched_location yield(parent_node, parent_field, node) end diff --git a/lib/pg_query/truncate.rb b/lib/pg_query/truncate.rb index 37cc1968..c7383ba8 100644 --- a/lib/pg_query/truncate.rb +++ b/lib/pg_query/truncate.rb @@ -60,7 +60,7 @@ def truncate(max_length) # rubocop:disable Metrics/CyclomaticComplexity def find_possible_truncations # rubocop:disable Metrics/CyclomaticComplexity truncations = [] - treewalker! @tree do |node, k, v, location| + treewalker_with_location! @tree do |node, k, v, location| case k when :target_list next unless node.is_a?(PgQuery::SelectStmt) || node.is_a?(PgQuery::UpdateStmt) || node.is_a?(PgQuery::OnConflictClause) diff --git a/spec/lib/treewalker_spec.rb b/spec/lib/treewalker_spec.rb index 3b7870c3..9784231c 100644 --- a/spec/lib/treewalker_spec.rb +++ b/spec/lib/treewalker_spec.rb @@ -1,6 +1,6 @@ require 'spec_helper' -describe PgQuery, '.treewalker' do +describe PgQuery, '#walk!' do it 'walks nodes contained in repeated fields' do locations = [] described_class.parse("SELECT to_timestamp($1)").walk! do |_, _, _, location|