Fix Code.Fragment.surround_context for keyword keys

lib/elixir/lib/code/fragment.ex

@@ -590,7 +590,8 @@ defmodule Code.Fragment do
                | {:sigil, charlist}
                | {:struct, inside_struct}
                | {:unquoted_atom, charlist}
-               | {:keyword, charlist},
+               | {:keyword, charlist}
+               | {:key, charlist},
              inside_dot:
                {:alias, charlist}
                | {:alias, inside_alias, charlist}
@@ -640,7 +641,16 @@ defmodule Code.Fragment do
         maybe_operator(reversed_pre, post, line, opts)
 
       {:identifier, reversed_post, rest} ->
-        {rest, _} = strip_spaces(rest, 0)
+        {keyword_key?, rest} =
+          case rest do
+            [?: | tail] when tail == [] or hd(tail) in @space ->
+              {true, rest}
+
+            _ ->
+              {rest, _} = strip_spaces(rest, 0)
+              {false, rest}
+          end
+
         reversed = reversed_post ++ reversed_pre
 
         case codepoint_cursor_context(reversed, opts) do
@@ -656,6 +666,9 @@ defmodule Code.Fragment do
           {{:dot, _, [_ | _]} = dot, offset} ->
             build_surround(dot, reversed, line, offset)
 
+          {{:local_or_var, acc}, offset} when keyword_key? ->
+            build_surround({:key, acc}, reversed, line, offset)
+
           {{:local_or_var, acc}, offset} when hd(rest) == ?( ->
             build_surround({:local_call, acc}, reversed, line, offset)
 

lib/elixir/test/elixir/code_fragment_test.exs

@@ -1110,6 +1110,44 @@ defmodule CodeFragmentTest do
 
       assert CF.surround_context(":", {1, 1}) == :none
     end
+
+    test "keyword keys" do
+      for i <- 2..4 do
+        assert CF.surround_context("[foo:", {1, i}) == %{
+                 context: {:key, ~c"foo"},
+                 begin: {1, 2},
+                 end: {1, 5}
+               }
+      end
+
+      for i <- 10..12 do
+        assert CF.surround_context("[foo: 1, bar: 2]", {1, i}) == %{
+                 context: {:key, ~c"bar"},
+                 begin: {1, 10},
+                 end: {1, 13}
+               }
+      end
+
+      assert CF.surround_context("if foo?, do: bar()", {1, 10}) == %{
+               context: {:key, ~c"do"},
+               begin: {1, 10},
+               end: {1, 12}
+             }
+    end
+
+    test "keyword false positives" do
+      assert CF.surround_context("<<foo::", {1, 3}) == %{
+               context: {:local_or_var, ~c"foo"},
+               begin: {1, 3},
+               end: {1, 6}
+             }
+
+      assert CF.surround_context("[foo  :atom", {1, 2}) == %{
+               context: {:local_or_var, ~c"foo"},
+               begin: {1, 2},
+               end: {1, 5}
+             }
+    end
   end
 
   describe "container_cursor_to_quoted/2" do