Improve create function code action (#1487)

The create function code action will now use the argument names used in the function call. It will also try to indent correctly by guessing indentation amount by analyzing the source file.
erlang-ls · Jan 19, 2024 · f901490 · f901490
1 parent 273c857
commit f901490
Show file tree

Hide file tree

Showing 4 changed files with 141 additions and 18 deletions.
diff --git a/apps/els_lsp/priv/code_navigation/src/code_action.erl b/apps/els_lsp/priv/code_navigation/src/code_action.erl
@@ -22,4 +22,5 @@ function_c() ->
 function_d() ->
   foobar(),
   foobar(x,y,z),
+  foobar(Foo, #foo_bar{}, Bar = 123, #foo_bar{} = Baz),
   ok.
diff --git a/apps/els_lsp/src/els_code_actions.erl b/apps/els_lsp/src/els_code_actions.erl
@@ -16,8 +16,10 @@
 
 -spec create_function(uri(), range(), binary(), [binary()]) -> [map()].
 create_function(Uri, Range0, _Data, [UndefinedFun]) ->
-    {ok, Document} = els_utils:lookup_document(Uri),
+    {ok, #{text := Text} = Document} = els_utils:lookup_document(Uri),
     Range = els_range:to_poi_range(Range0),
+    Indent = guess_indentation(string:lexemes(Text, "\n")),
+    IndentStr = lists:duplicate(Indent, 32),
     FunPOIs = els_dt_document:pois(Document, [function]),
     %% Figure out which function the error was found in, as we want to
     %% create the function right after the current function.
@@ -32,8 +34,11 @@ create_function(Uri, Range0, _Data, [UndefinedFun]) ->
         [#{to := {Line, _}} | _] ->
             [Name, ArityBin] = string:split(UndefinedFun, "/"),
             Arity = binary_to_integer(ArityBin),
-            Args = string:join(lists:duplicate(Arity, "_"), ", "),
-            SpecAndFun = io_lib:format("~s(~s) ->\n  ok.\n\n", [Name, Args]),
+            Args = format_args(Document, Arity, Range),
+            SpecAndFun = io_lib:format(
+                "~s(~s) ->\n~sok.\n\n",
+                [Name, Args, IndentStr]
+            ),
             [
                 make_edit_action(
                     Uri,
@@ -244,3 +249,45 @@ make_edit_action(Uri, Title, Kind, Text, Range) ->
 -spec edit(uri(), binary(), range()) -> workspace_edit().
 edit(Uri, Text, Range) ->
     #{changes => #{Uri => [#{newText => Text, range => Range}]}}.
+
+-spec format_args(
+    els_dt_document:item(),
+    non_neg_integer(),
+    els_poi:poi_range()
+) -> string().
+format_args(Document, Arity, Range) ->
+    %% Find the matching function application and extract
+    %% argument names from it.
+    AppPOIs = els_dt_document:pois(Document, [application]),
+    Matches = [
+        POI
+     || #{range := R} = POI <- AppPOIs,
+        els_range:in(R, Range)
+    ],
+    case Matches of
+        [#{data := #{args := Args0}} | _] ->
+            string:join([A || {_N, A} <- Args0], ", ");
+        [] ->
+            string:join(lists:duplicate(Arity, "_"), ", ")
+    end.
+
+-spec guess_indentation([binary()]) -> pos_integer().
+guess_indentation([]) ->
+    2;
+guess_indentation([A, B | Rest]) ->
+    ACount = count_leading_spaces(A, 0),
+    BCount = count_leading_spaces(B, 0),
+    case {ACount, BCount} of
+        {0, N} when N > 0 ->
+            N;
+        {_, _} ->
+            guess_indentation([B | Rest])
+    end.
+
+-spec count_leading_spaces(binary(), non_neg_integer()) -> non_neg_integer().
+count_leading_spaces(<<>>, _Acc) ->
+    0;
+count_leading_spaces(<<" ", Rest/binary>>, Acc) ->
+    count_leading_spaces(Rest, 1 + Acc);
+count_leading_spaces(<<_:8, _/binary>>, Acc) ->
+    Acc.
diff --git a/apps/els_lsp/src/els_parser.erl b/apps/els_lsp/src/els_parser.erl
@@ -290,9 +290,12 @@ application(Tree) ->
             Pos = erl_syntax:get_pos(erl_syntax:application_operator(Tree)),
             case erl_internal:bif(F, A) of
                 %% Call to a function from the `erlang` module
-                true -> [poi(Pos, application, {erlang, F, A}, #{imported => true})];
+                true ->
+                    [poi(Pos, application, {erlang, F, A}, #{imported => true})];
                 %% Local call
-                false -> [poi(Pos, application, {F, A})]
+                false ->
+                    Args = erl_syntax:application_arguments(Tree),
+                    [poi(Pos, application, {F, A}, #{args => args_from_subtrees(Args)})]
             end;
         {{ModType, M}, {FunType, F}, A} ->
             ModFunTree = erl_syntax:application_operator(Tree),
@@ -719,14 +722,44 @@ function_args(Clause) ->
 args_from_subtrees(Trees) ->
     Arity = length(Trees),
     [
-        case erl_syntax:type(T) of
-            %% TODO: Handle literals
-            variable -> {N, erl_syntax:variable_literal(T)};
-            _ -> {N, "Arg" ++ integer_to_list(N)}
+        case extract_variable(T) of
+            {true, Variable} ->
+                {N, Variable};
+            false ->
+                {N, "Arg" ++ integer_to_list(N)}
         end
      || {N, T} <- lists:zip(lists:seq(1, Arity), Trees)
     ].
 
+-spec extract_variable(tree()) -> {true, string()} | false.
+extract_variable(T) ->
+    case erl_syntax:type(T) of
+        %% TODO: Handle literals
+        variable ->
+            {true, erl_syntax:variable_literal(T)};
+        match_expr ->
+            Body = erl_syntax:match_expr_body(T),
+            Pattern = erl_syntax:match_expr_pattern(T),
+            case {extract_variable(Pattern), extract_variable(Body)} of
+                {false, Result} ->
+                    Result;
+                {Result, _} ->
+                    Result
+            end;
+        record_expr ->
+            RecordNode = erl_syntax:record_expr_type(T),
+            case erl_syntax:type(RecordNode) of
+                atom ->
+                    NameAtom = erl_syntax:atom_value(RecordNode),
+                    NameBin = els_utils:camel_case(atom_to_binary(NameAtom, utf8)),
+                    {true, unicode:characters_to_list(NameBin)};
+                _ ->
+                    false
+            end;
+        _Type ->
+            false
+    end.
+
 -spec implicit_fun(tree()) -> [els_poi:poi()].
 implicit_fun(Tree) ->
     FunSpec =

diff --git a/apps/els_lsp/test/els_code_action_SUITE.erl b/apps/els_lsp/test/els_code_action_SUITE.erl
@@ -20,6 +20,7 @@
     remove_unused_import/1,
     create_undefined_function/1,
     create_undefined_function_arity/1,
+    create_undefined_function_variable_names/1,
     fix_callbacks/1
 ]).
 
@@ -313,8 +314,8 @@ remove_unused_import(Config) ->
 create_undefined_function(Config) ->
     Uri = ?config(code_action_uri, Config),
     Range = els_protocol:range(#{
-        from => {23, 2},
-        to => {23, 8}
+        from => {23, 3},
+        to => {23, 9}
     }),
     Diag = #{
         message => <<"function foobar/0 undefined">>,
@@ -334,8 +335,8 @@ create_undefined_function(Config) ->
                                     #{
                                         range =>
                                             els_protocol:range(#{
-                                                from => {27, 1},
-                                                to => {27, 1}
+                                                from => {28, 1},
+                                                to => {28, 1}
                                             }),
                                         newText =>
                                             <<"foobar() ->\n  ok.\n\n">>
@@ -354,8 +355,8 @@ create_undefined_function(Config) ->
 create_undefined_function_arity(Config) ->
     Uri = ?config(code_action_uri, Config),
     Range = els_protocol:range(#{
-        from => {24, 2},
-        to => {24, 8}
+        from => {24, 3},
+        to => {24, 9}
     }),
     Diag = #{
         message => <<"function foobar/3 undefined">>,
@@ -375,11 +376,11 @@ create_undefined_function_arity(Config) ->
                                     #{
                                         range =>
                                             els_protocol:range(#{
-                                                from => {27, 1},
-                                                to => {27, 1}
+                                                from => {28, 1},
+                                                to => {28, 1}
                                             }),
                                         newText =>
-                                            <<"foobar(_, _, _) ->\n  ok.\n\n">>
+                                            <<"foobar(Arg1, Arg2, Arg3) ->\n  ok.\n\n">>
                                     }
                                 ]
                         }
@@ -391,6 +392,47 @@ create_undefined_function_arity(Config) ->
     ?assertEqual(Expected, Result),
     ok.
 
+-spec create_undefined_function_variable_names((config())) -> ok.
+create_undefined_function_variable_names(Config) ->
+    Uri = ?config(code_action_uri, Config),
+    Range = els_protocol:range(#{
+        from => {25, 3},
+        to => {25, 9}
+    }),
+    Diag = #{
+        message => <<"function foobar/5 undefined">>,
+        range => Range,
+        severity => 2,
+        source => <<"">>
+    },
+    #{result := Result} = els_client:document_codeaction(Uri, Range, [Diag]),
+    Expected =
+        [
+            #{
+                edit => #{
+                    changes =>
+                        #{
+                            binary_to_atom(Uri, utf8) =>
+                                [
+                                    #{
+                                        range =>
+                                            els_protocol:range(#{
+                                                from => {28, 1},
+                                                to => {28, 1}
+                                            }),
+                                        newText =>
+                                            <<"foobar(Foo, FooBar, Bar, FooBar) ->\n  ok.\n\n">>
+                                    }
+                                ]
+                        }
+                },
+                kind => <<"quickfix">>,
+                title => <<"Create function foobar/5">>
+            }
+        ],
+    ?assertEqual(Expected, Result),
+    ok.
+
 -spec fix_callbacks(config()) -> ok.
 fix_callbacks(Config) ->
     Uri = ?config(code_action_uri, Config),