Support node groups in Node#each_descendant and similar traversal m…

…ethods

Sometimes it would be nice to do `each_child_node(:call)` instead of `each_child_node(:send, :csend)`.
This nicely mirrors that you can already use these aliases when writing node patterns

changelog/change_allow_groups_for_traversal.md

@@ -0,0 +1 @@
+* [#357](https://github.com/rubocop/rubocop-ast/pull/357): Support node groups in `Node#each_descendant` and similar traversal methods. ([@earlopain][])

lib/rubocop/ast/node.rb

@@ -76,9 +76,6 @@ class Node < Parser::AST::Node # rubocop:disable Metrics/ClassLength
       OPERATOR_KEYWORDS = %i[and or].to_set.freeze
       # @api private
       SPECIAL_KEYWORDS = %w[__FILE__ __LINE__ __ENCODING__].to_set.freeze
-      # @api private
-      ARGUMENT_TYPES = %i[arg optarg restarg kwarg kwoptarg kwrestarg
-                          blockarg forward_arg shadowarg].to_set.freeze
 
       LITERAL_RECURSIVE_METHODS = (COMPARISON_OPERATORS + %i[* ! <=>]).freeze
       LITERAL_RECURSIVE_TYPES = (OPERATOR_KEYWORDS + COMPOSITE_LITERALS + %i[begin pair]).freeze
@@ -98,7 +95,7 @@ class Node < Parser::AST::Node # rubocop:disable Metrics/ClassLength
         kwrestarg: :argument,
         blockarg: :argument,
         forward_arg: :argument,
-        shardowarg: :argument,
+        shadowarg: :argument,
 
         true: :boolean,
         false: :boolean,
@@ -657,8 +654,7 @@ def visit_ancestors(types)
         last_node = self
 
         while (current_node = last_node.parent)
-          yield current_node if types.empty? ||
-                                types.include?(current_node.type)
+          yield current_node if types.empty? || current_node.type?(*types)
           last_node = current_node
         end
       end

lib/rubocop/ast/node/args_node.rb

@@ -32,7 +32,7 @@ def empty_and_without_delimiters?
       #
       # @return [Array<Node>] array of argument nodes.
       def argument_list
-        each_descendant(*ARGUMENT_TYPES).to_a.freeze
+        each_descendant(:argument).to_a.freeze
       end
     end
   end

lib/rubocop/ast/node/mixin/descendence.rb

@@ -25,7 +25,7 @@ def each_child_node(*types)
         children.each do |child|
           next unless child.is_a?(::AST::Node)
 
-          yield child if types.empty? || types.include?(child.type)
+          yield child if types.empty? || child.type?(*types)
         end
 
         self
@@ -95,7 +95,7 @@ def descendants
       def each_node(*types, &block)
         return to_enum(__method__, *types) unless block
 
-        yield self if types.empty? || types.include?(type)
+        yield self if types.empty? || type?(*types)
 
         visit_descendants(types, &block)
 
@@ -108,7 +108,7 @@ def visit_descendants(types, &block)
         children.each do |child|
           next unless child.is_a?(::AST::Node)
 
-          yield child if types.empty? || types.include?(child.type)
+          yield child if types.empty? || child.type?(*types)
           child.visit_descendants(types, &block)
         end
       end

spec/rubocop/ast/node_spec.rb

@@ -1133,4 +1133,52 @@ class << expr
       expect(node).not_to be_loc_is(:foo, ':')
     end
   end
+
+  context 'traversal' do
+    shared_examples 'traverse' do |test_name, *types, result|
+      it "handles #{test_name}" do
+        expect(node.send(method, *types).map(&:type)).to eq(result)
+      end
+    end
+
+    describe '#each_ancestor' do
+      let(:src) { 'foo&.bar([>>baz<<]).bat' }
+      let(:method) { :each_ancestor }
+
+      it_behaves_like 'traverse', 'no argument', %i[array csend send]
+      it_behaves_like 'traverse', 'single argument', :send, %i[send]
+      it_behaves_like 'traverse', 'two arguments', :send, :csend, %i[csend send]
+      it_behaves_like 'traverse', 'group argument', :call, %i[csend send]
+    end
+
+    describe '#each_child_node' do
+      let(:src) { '[1, 2.0, foo, 3i, 4, BAR, 5r]' }
+      let(:method) { :each_child_node }
+
+      it_behaves_like 'traverse', 'no argument', %i[int float send complex int const rational]
+      it_behaves_like 'traverse', 'single argument', :int, %i[int int]
+      it_behaves_like 'traverse', 'two arguments', :int, :complex, :csend, %i[int complex int]
+      it_behaves_like 'traverse', 'group argument', :numeric, %i[int float complex int rational]
+    end
+
+    describe '#each_descendant' do
+      let(:src) { 'foo(true, false, bar(baz, true))' }
+      let(:method) { :each_descendant }
+
+      it_behaves_like 'traverse', 'no argument', %i[true false send send true]
+      it_behaves_like 'traverse', 'single argument', :true, %i[true true]
+      it_behaves_like 'traverse', 'two arguments', :false, :send, %i[false send send]
+      it_behaves_like 'traverse', 'group argument', :boolean, %i[true false true]
+    end
+
+    describe '#each_node' do
+      let(:src) { 'def foo(bar, *, baz, **kw, &block); 123; end' }
+      let(:method) { :each_node }
+
+      it_behaves_like 'traverse', 'no argument', %i[def args arg restarg arg kwrestarg blockarg int]
+      it_behaves_like 'traverse', 'single argument', :arg, %i[arg arg]
+      it_behaves_like 'traverse', 'two arguments', :restarg, :blockarg, %i[restarg blockarg]
+      it_behaves_like 'traverse', 'group argument', :argument, %i[arg restarg arg kwrestarg blockarg]
+    end
+  end
 end