NodePattern: Optimize unions of literals into Sets

lib/rubocop/ast.rb

@@ -20,6 +20,7 @@
 require_relative 'ast/node_pattern/lexer'
 require_relative 'ast/node_pattern/node'
 require_relative 'ast/node_pattern/parser'
+require_relative 'ast/node_pattern/sets'
 require_relative 'ast/sexp'
 require_relative 'ast/node'
 require_relative 'ast/node/mixin/method_identifier_predicates'

lib/rubocop/ast/node_pattern/builder.rb

@@ -40,7 +40,8 @@ def emit_union(begin_t, pattern_lists, end_t)
               emit_subsequence(list)
             end
           end
-          emit_list(:union, begin_t, children, end_t)
+          type = set_optimizable?(children) ? :set : :union
+          emit_list(type, begin_t, children, end_t)
         end
 
         def emit_subsequence(node_list)
@@ -51,6 +52,10 @@ def emit_subsequence(node_list)
 
         private
 
+        def set_optimizable?(children)
+          children.all?(&:literal?)
+        end
+
         def n(type, *args)
           Node::MAP[type].new(type, *args)
         end

lib/rubocop/ast/node_pattern/compiler/atom_subcompiler.rb

@@ -36,6 +36,11 @@ def visit_positional_parameter
             compiler.positional_parameter(node.child)
           end
 
+          def visit_set
+            set = node.children.map(&:child).to_set.freeze
+            NodePattern::Sets[set]
+          end
+
           # Assumes other types are node patterns.
           def visit_other_type
             compiler.with_temp_variables do |compare|

lib/rubocop/ast/node_pattern/node.rb

@@ -8,6 +8,9 @@ class Node < ::Parser::AST::Node
         extend Forwardable
         include ::RuboCop::AST::Descendence
 
+        LITERAL_TYPES = %i[symbol number string].to_set.freeze
+        private_constant :LITERAL_TYPES
+
         ###
         # To be overriden by subclasses
         ###
@@ -55,6 +58,11 @@ def variadic?
           arity.is_a?(Range)
         end
 
+        # @return [Boolean] returns true for nodes having a Ruby literal equivalent
+        def literal?
+          LITERAL_TYPES.include?(type)
+        end
+
         # @return [Range] arity as a Range
         def arity_range
           a = arity

lib/rubocop/ast/node_pattern/sets.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module RuboCop
+  module AST
+    class NodePattern
+      # Utility to assign a set of values to a constant
+      module Sets
+        REGISTRY = Hash.new do |h, set|
+          name = Sets.name(set).freeze
+          Sets.const_set(name, set)
+          h[set] = "::RuboCop::AST::NodePattern::Sets::#{name}"
+        end
+
+        MAX = 4
+        def self.name(set)
+          elements = set
+          elements = set.first(MAX - 1) << :etc if set.size > MAX
+          name = elements.to_a.join('_').upcase.gsub(/[^A-Z0-9_]/, '')
+          uniq("SET_#{name}")
+        end
+
+        def self.uniq(name)
+          return name unless Sets.const_defined?(name)
+
+          (2..Float::INFINITY).each do |i|
+            uniq = "#{name}_#{i}"
+            return uniq unless Sets.const_defined?(uniq)
+          end
+        end
+
+        def self.[](set)
+          REGISTRY[set]
+        end
+      end
+    end
+  end
+end

spec/rubocop/ast/node_pattern/parser_spec.rb

@@ -61,6 +61,16 @@
       )
     end
 
+    it 'parses unions of literals as a set' do
+      expect_parsing(
+        s(:sequence, s(:set, s(:symbol, :a), s(:number, 42), s(:string, 'hello'))),
+        '({:a 42 "hello"})',
+        ' ^               begin (set)
+        |               ^ end (set)
+        |     ~~          expression (set/1.number)'
+      )
+    end
+
     it 'generates specialized nodes' do
       source_file = Parser::Source::Buffer.new('(spec)', source: '($_)')
       ast = parser.parse(source_file)

spec/rubocop/ast/node_pattern/sets_spec.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+RSpec.describe RuboCop::AST::NodePattern::Sets do
+  subject(:name) { described_class[set] }
+
+  let(:set) { Set[1, 2, 3, 4, 5, 6] }
+
+  it { is_expected.to eq '::RuboCop::AST::NodePattern::Sets::SET_1_2_3_ETC' }
+
+  it { is_expected.to eq described_class[Set[6, 5, 4, 3, 2, 1]] }
+
+  it { is_expected.not_to eq described_class[Set[1, 2, 3, 4, 5, 6, 7]] }
+
+  it 'creates a constant with the right value' do
+    expect(eval(name)).to eq set # rubocop:disable Security/Eval
+  end
+end