From a9f7be9e360732077cea1eb3e323e0fcac38b712 Mon Sep 17 00:00:00 2001
From: Earlopain <14981592+Earlopain@users.noreply.github.com>
Date: Wed, 22 Jan 2025 21:24:11 +0100
Subject: [PATCH] [Fix #323] Fix node captures inside of `?`, `+`, and `*`
 repetition

The calculation for how many captures a pattern has was not correct in this case.
A pattern like `(send _ _ (int {$_ | $_ | $_})?)` would end up with 3 captures, even though
the union would in total only contribute one
---
 changelog/fix_captures_in_repetition.md |   1 +
 lib/rubocop/ast/node_pattern/node.rb    |   8 +-
 spec/rubocop/ast/node_pattern_spec.rb   | 162 ++++++++++++++++++++++++
 3 files changed, 170 insertions(+), 1 deletion(-)
 create mode 100644 changelog/fix_captures_in_repetition.md

diff --git a/changelog/fix_captures_in_repetition.md b/changelog/fix_captures_in_repetition.md
new file mode 100644
index 000000000..fefccd17e
--- /dev/null
+++ b/changelog/fix_captures_in_repetition.md
@@ -0,0 +1 @@
+* [#323](https://github.com/rubocop/rubocop-ast/issues/323): Fix node captures inside of `?`, `+`, and `*` repetition. ([@earlopain][])
diff --git a/lib/rubocop/ast/node_pattern/node.rb b/lib/rubocop/ast/node_pattern/node.rb
index 8645f9f07..95ff3ccd8 100644
--- a/lib/rubocop/ast/node_pattern/node.rb
+++ b/lib/rubocop/ast/node_pattern/node.rb
@@ -48,7 +48,7 @@ def child
           children[0]
         end
 
-        # @return [Integer] nb of captures of that node and its descendants
+        # @return [Integer] nb of captures that this node will emit
         def nb_captures
           children_nodes.sum(&:nb_captures)
         end
@@ -243,6 +243,12 @@ def in_sequence_head
 
             [with(children: new_children)]
           end
+
+          # Each child in a union must contain the same number
+          # of captures. Only one branch ends up capturing.
+          def nb_captures
+            child.nb_captures
+          end
         end
 
         # Registry
diff --git a/spec/rubocop/ast/node_pattern_spec.rb b/spec/rubocop/ast/node_pattern_spec.rb
index 5d65f7fde..fb01e0389 100644
--- a/spec/rubocop/ast/node_pattern_spec.rb
+++ b/spec/rubocop/ast/node_pattern_spec.rb
@@ -958,6 +958,168 @@ def not_match_codes(*codes)
 
       it_behaves_like 'invalid'
     end
+
+    context 'using repetition' do
+      shared_examples 'repetition' do |behavior|
+        context 'with one capture' do
+          let(:pattern) { pattern_placeholder.sub('{}', '{$_}') }
+
+          it_behaves_like behavior
+        end
+
+        context 'with two captures' do
+          let(:pattern) { pattern_placeholder.sub('{}', '{$_ | $_}') }
+
+          it_behaves_like behavior
+        end
+
+        context 'with three captures' do
+          let(:pattern) { pattern_placeholder.sub('{}', '{$_ | $_ | $_}') }
+
+          it_behaves_like behavior
+        end
+      end
+
+      context 'with ?' do
+        context 'one capture' do
+          let(:pattern_placeholder) { '(send _ _ (int {})?)' }
+
+          context 'one match' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo(1)' }
+              let(:captured_val) { [1] }
+            end
+          end
+
+          context 'no match' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo' }
+              let(:captured_val) { [] }
+            end
+          end
+        end
+
+        context 'two captures' do
+          let(:pattern_placeholder) { '(send _ $_ (int {})?)' }
+
+          context 'one match' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo(1)' }
+              let(:captured_vals) { [:foo, [1]] }
+            end
+          end
+
+          context 'no match' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo' }
+              let(:captured_vals) { [:foo, []] }
+            end
+          end
+        end
+      end
+
+      context 'with +' do
+        context 'one capture' do
+          let(:pattern_placeholder) { '(send _ _ (int {})+)' }
+
+          context 'one match' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo(1)' }
+              let(:captured_val) { [1] }
+            end
+          end
+
+          context 'two matches' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo(1, 2)' }
+              let(:captured_val) { [1, 2] }
+            end
+          end
+
+          context 'no match' do
+            it_behaves_like 'repetition', 'nonmatching' do
+              let(:ruby) { 'foo' }
+            end
+          end
+        end
+
+        context 'two captures' do
+          let(:pattern_placeholder) { '(send _ $_ (int {})+)' }
+
+          context 'one match' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo(1)' }
+              let(:captured_vals) { [:foo, [1]] }
+            end
+          end
+
+          context 'two matches' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo(1, 2)' }
+              let(:captured_vals) { [:foo, [1, 2]] }
+            end
+          end
+
+          context 'no match' do
+            it_behaves_like 'repetition', 'nonmatching' do
+              let(:ruby) { 'foo' }
+            end
+          end
+        end
+      end
+
+      context 'with *' do
+        context 'one capture' do
+          let(:pattern_placeholder) { '(send _ _ (int {})*)' }
+
+          context 'one match' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo(1)' }
+              let(:captured_val) { [1] }
+            end
+          end
+
+          context 'two matches' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo(1, 2)' }
+              let(:captured_val) { [1, 2] }
+            end
+          end
+
+          context 'no match' do
+            it_behaves_like 'repetition', 'single capture' do
+              let(:ruby) { 'foo' }
+              let(:captured_val) { [] }
+            end
+          end
+        end
+
+        context 'two captures' do
+          let(:pattern_placeholder) { '(send _ $_ (int {})*)' }
+
+          context 'one match' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo(1)' }
+              let(:captured_vals) { [:foo, [1]] }
+            end
+          end
+
+          context 'two matches' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo(1, 2)' }
+              let(:captured_vals) { [:foo, [1, 2]] }
+            end
+          end
+
+          context 'no match' do
+            it_behaves_like 'repetition', 'multiple capture' do
+              let(:ruby) { 'foo' }
+              let(:captured_vals) { [:foo, []] }
+            end
+          end
+        end
+      end
+    end
   end
 
   describe 'negation' do