From 1da77c8648a61b42fe479a4e2aff38dbb1841cdd Mon Sep 17 00:00:00 2001
From: sabiwara <sabiwara@gmail.com>
Date: Sat, 21 Sep 2024 15:22:39 +0900
Subject: [PATCH] Code.Fragment.surround_context/2 handles &123 as :capture_arg

---
 lib/elixir/lib/code/fragment.ex               | 27 +++++++++++++++-
 lib/elixir/test/elixir/code_fragment_test.exs | 32 +++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/lib/elixir/lib/code/fragment.ex b/lib/elixir/lib/code/fragment.ex
index d759cdaa687..c78b437dfd4 100644
--- a/lib/elixir/lib/code/fragment.ex
+++ b/lib/elixir/lib/code/fragment.ex
@@ -320,6 +320,9 @@ defmodule Code.Fragment do
           _ ->
             {{:local_or_var, acc}, count}
         end
+
+      {:capture_arg, acc, count} ->
+        {{:capture_arg, acc}, count}
     end
   end
 
@@ -362,6 +365,14 @@ defmodule Code.Fragment do
     :none
   end
 
+  defp rest_identifier([?& | tail] = rest, count, acc) when tail == [] or hd(tail) != ?& do
+    if Enum.all?(acc, &(&1 in ?0..?9)) do
+      {:capture_arg, [?& | acc], count + 1}
+    else
+      tokenize_identifier(rest, count, acc)
+    end
+  end
+
   defp rest_identifier(rest, count, acc) do
     tokenize_identifier(rest, count, acc)
   end
@@ -591,7 +602,8 @@ defmodule Code.Fragment do
                | {:struct, inside_struct}
                | {:unquoted_atom, charlist}
                | {:keyword, charlist}
-               | {:key, charlist},
+               | {:key, charlist}
+               | {:capture_arg, charlist},
              inside_dot:
                {:alias, charlist}
                | {:alias, inside_alias, charlist}
@@ -696,6 +708,9 @@ defmodule Code.Fragment do
           {{:unquoted_atom, acc}, offset} ->
             build_surround({:unquoted_atom, acc}, reversed, line, offset)
 
+          {{:capture_arg, acc}, offset} ->
+            build_surround({:capture_arg, acc}, reversed, line, offset)
+
           _ ->
             maybe_operator(reversed_pre, post, line, opts)
         end
@@ -728,6 +743,16 @@ defmodule Code.Fragment do
         reversed = reversed_post ++ reversed_pre
 
         case codepoint_cursor_context(reversed, opts) do
+          {{:operator, ~c"&"}, offset} when hd(rest) in ?0..?9 ->
+            arg = Enum.take_while(rest, &(&1 in ?0..?9))
+
+            build_surround(
+              {:capture_arg, ~c"&" ++ arg},
+              :lists.reverse(arg, reversed),
+              line,
+              offset + length(arg)
+            )
+
           {{:operator, acc}, offset} ->
             build_surround({:operator, acc}, reversed, line, offset)
 
diff --git a/lib/elixir/test/elixir/code_fragment_test.exs b/lib/elixir/test/elixir/code_fragment_test.exs
index 37937b87fe6..3ceb4765395 100644
--- a/lib/elixir/test/elixir/code_fragment_test.exs
+++ b/lib/elixir/test/elixir/code_fragment_test.exs
@@ -1091,6 +1091,38 @@ defmodule CodeFragmentTest do
       end
     end
 
+    test "capture operator" do
+      assert CF.surround_context("& &123 + 1", {1, 1}) == %{
+               context: {:operator, ~c"&"},
+               begin: {1, 1},
+               end: {1, 2}
+             }
+
+      for i <- 3..6 do
+        assert CF.surround_context("& &123 + 1", {1, i}) == %{
+                 context: {:capture_arg, ~c"&123"},
+                 begin: {1, 3},
+                 end: {1, 7}
+               }
+      end
+    end
+
+    test "capture operator false positive" do
+      assert CF.surround_context("1&&2", {1, 3}) == %{
+               context: {:operator, ~c"&&"},
+               begin: {1, 2},
+               end: {1, 4}
+             }
+
+      assert CF.surround_context("1&&2", {1, 4}) == :none
+
+      assert CF.surround_context("&a", {1, 2}) == %{
+               context: {:local_or_var, ~c"a"},
+               begin: {1, 2},
+               end: {1, 3}
+             }
+    end
+
     test "unquoted atom" do
       for i <- 1..10 do
         assert CF.surround_context(":hello_wor", {1, i}) == %{