[erlang-ls#1337] Make goto variable definition respect list comprehen…

…sion scopes
plux · Dec 20, 2023 · 353e820 · 353e820
1 parent cee4096
commit 353e820
Show file tree

Hide file tree

Showing 5 changed files with 256 additions and 27 deletions.
diff --git a/apps/els_lsp/priv/code_navigation/src/variable_list_comp.erl b/apps/els_lsp/priv/code_navigation/src/variable_list_comp.erl
@@ -0,0 +1,20 @@
+-module(variable_list_comp).
+
+one() ->
+    Var = 1,
+    [ Var || Var <- [1, 2, 3] ],
+    Var.
+
+two() ->
+    [ Var || Var <- [1, 2, 3] ],
+    [ Var || Var <- [4, 5, 6] ].
+
+three() ->
+    Var = 1,
+    [ Var || _ <- [1, 2, 3] ],
+    Var.
+
+four() ->
+    [ {Var, Var2} || Var <- [4, 5, 6],
+                     Var2 <- [ Var || Var <- [1, 2, 3] ]
+    ].
diff --git a/apps/els_lsp/src/els_code_navigation.erl b/apps/els_lsp/src/els_code_navigation.erl
@@ -271,13 +271,63 @@ maybe_imported(_Document, _Kind, _Data) ->
     [].
 
 -spec find_in_scope(uri(), els_poi:poi()) -> [els_poi:poi()].
-find_in_scope(Uri, #{kind := variable, id := VarId, range := VarRange}) ->
+find_in_scope(
+    Uri,
+    #{kind := variable, id := VarId, range := VarRange}
+) ->
     {ok, Document} = els_utils:lookup_document(Uri),
+    LcPOIs = els_poi:sort(els_dt_document:pois(Document, [list_comp])),
     VarPOIs = els_poi:sort(els_dt_document:pois(Document, [variable])),
     ScopeRange = els_scope:variable_scope_range(VarRange, Document),
-    [
+    MatchInScope = [
         POI
-     || #{range := Range, id := Id} = POI <- VarPOIs,
-        els_range:in(Range, ScopeRange),
+     || #{id := Id} = POI <- pois_in(VarPOIs, ScopeRange),
         Id =:= VarId
-    ].
+    ],
+    %% Handle special case if variable POI is inside list comprehension (LC)
+    MatchingLcPOIs = pois_contains(LcPOIs, VarRange),
+    case find_in_scope_list_comp(MatchingLcPOIs, MatchInScope) of
+        [] ->
+            MatchInScope -- find_in_scope_list_comp(LcPOIs, MatchInScope);
+        MatchInLc ->
+            MatchInLc
+    end.
+
+-spec find_in_scope_list_comp([els_poi:poi()], [els_poi:poi()]) ->
+    [els_poi:poi()].
+find_in_scope_list_comp([], _VarPOIs) ->
+    %% No match in LC, use regular scope
+    [];
+find_in_scope_list_comp([LcPOI | LcPOIs], VarPOIs) ->
+    #{data := #{pattern_ranges := PatRanges}, range := LcRange} = LcPOI,
+    VarsInLc = pois_in(VarPOIs, LcRange),
+    {PatVars, OtherVars} =
+        lists:partition(
+            fun(#{range := Range}) ->
+                lists:any(
+                    fun(PatRange) ->
+                        els_range:in(Range, PatRange)
+                    end,
+                    PatRanges
+                )
+            end,
+            VarsInLc
+        ),
+    case PatVars of
+        [] ->
+            %% Didn't find any patterned vars in this LC, try the next one
+            find_in_scope_list_comp(LcPOIs, VarPOIs);
+        _ ->
+            %% Put pattern vars first to make goto definition work
+            PatVars ++ OtherVars
+    end.
+
+-spec pois_in([els_poi:poi()], els_poi:poi_range()) ->
+    [els_poi:poi()].
+pois_in(POIs, Range) ->
+    [POI || #{range := R} = POI <- POIs, els_range:in(R, Range)].
+
+-spec pois_contains([els_poi:poi()], els_poi:poi_range()) ->
+    [els_poi:poi()].
+pois_contains(POIs, Range) ->
+    [POI || #{range := R} = POI <- POIs, els_range:in(Range, R)].
diff --git a/apps/els_lsp/src/els_parser.erl b/apps/els_lsp/src/els_parser.erl
@@ -213,6 +213,8 @@ do_points_of_interest(Tree) ->
                 record_index_expr(Tree);
             record_expr ->
                 record_expr(Tree);
+            list_comp ->
+                list_comp(Tree);
             variable ->
                 variable(Tree);
             atom ->
@@ -721,6 +723,21 @@ macro(Tree) ->
     Anno = macro_location(Tree),
     [poi(Anno, macro, macro_name(Tree))].
 
+-spec list_comp(tree()) -> [els_poi:poi()].
+list_comp(Tree) ->
+    Pos = erl_syntax:get_pos(Tree),
+    Body = erl_syntax:list_comp_body(Tree),
+    PatRanges = [
+        els_range:range(
+            erl_syntax:get_pos(
+                erl_syntax:generator_pattern(Gen)
+            )
+        )
+     || Gen <- Body,
+        erl_syntax:type(Gen) =:= generator
+    ],
+    [poi(Pos, list_comp, undefined, #{pattern_ranges => PatRanges})].
+
 -spec map_record_def_fields(Fun, tree(), atom()) -> [Result] when
     Fun :: fun((tree(), atom()) -> Result).
 map_record_def_fields(Fun, Fields, RecordName) ->

diff --git a/apps/els_lsp/test/els_definition_SUITE.erl b/apps/els_lsp/test/els_definition_SUITE.erl
@@ -53,6 +53,7 @@
     type_application_user/1,
     type_export_entry/1,
     variable/1,
+    variable_list_comp/1,
     opaque_application_remote/1,
     opaque_application_user/1,
     parse_incomplete/1
@@ -623,38 +624,77 @@ type_export_entry(Config) ->
 -spec variable(config()) -> ok.
 variable(Config) ->
     Uri = ?config(code_navigation_uri, Config),
-    Def0 = els_client:definition(Uri, 104, 9),
-    Def1 = els_client:definition(Uri, 105, 10),
-    Def2 = els_client:definition(Uri, 107, 10),
-    Def3 = els_client:definition(Uri, 108, 10),
-    Def4 = els_client:definition(Uri, 19, 36),
-    #{result := [#{range := Range0, uri := DefUri0}]} = Def0,
-    #{result := [#{range := Range1, uri := DefUri0}]} = Def1,
-    #{result := [#{range := Range2, uri := DefUri0}]} = Def2,
-    #{result := [#{range := Range3, uri := DefUri0}]} = Def3,
-    #{result := [#{range := Range4, uri := DefUri0}]} = Def4,
-
-    ?assertEqual(?config(code_navigation_uri, Config), DefUri0),
     ?assertEqual(
-        els_protocol:range(#{from => {103, 12}, to => {103, 15}}),
-        Range0
+        {#{from => {103, 12}, to => {103, 15}}, Uri},
+        definition(Uri, 104, 9)
     ),
     ?assertEqual(
-        els_protocol:range(#{from => {104, 3}, to => {104, 6}}),
-        Range1
+        {#{from => {104, 3}, to => {104, 6}}, Uri},
+        definition(Uri, 105, 10)
     ),
     ?assertEqual(
-        els_protocol:range(#{from => {106, 12}, to => {106, 15}}),
-        Range2
+        {#{from => {106, 12}, to => {106, 15}}, Uri},
+        definition(Uri, 107, 10)
     ),
     ?assertEqual(
-        els_protocol:range(#{from => {106, 12}, to => {106, 15}}),
-        Range3
+        {#{from => {106, 12}, to => {106, 15}}, Uri},
+        definition(Uri, 108, 10)
     ),
     %% Inside macro
     ?assertEqual(
-        els_protocol:range(#{from => {19, 17}, to => {19, 18}}),
-        Range4
+        {#{from => {19, 17}, to => {19, 18}}, Uri},
+        definition(Uri, 19, 36)
+    ),
+    ok.
+
+-spec variable_list_comp(config()) -> ok.
+variable_list_comp(Config) ->
+    Uri = ?config(variable_list_comp_uri, Config),
+
+    %% one: Should skip LC
+    ?assertEqual(
+        {#{from => {4, 5}, to => {4, 8}}, Uri},
+        definition(Uri, 6, 5)
+    ),
+    %% one: Should go to LC generator pattern
+    ?assertEqual(
+        {#{from => {5, 14}, to => {5, 17}}, Uri},
+        definition(Uri, 5, 7)
+    ),
+    %% two: Should go to first LC generator pattern
+    ?assertEqual(
+        {#{from => {9, 14}, to => {9, 17}}, Uri},
+        definition(Uri, 9, 7)
+    ),
+    %% two: Should go to second LC generator pattern
+    ?assertEqual(
+        {#{from => {10, 14}, to => {10, 17}}, Uri},
+        definition(Uri, 10, 7)
+    ),
+    %% three: Should go to variable definition
+    ?assertEqual(
+        {#{from => {13, 5}, to => {13, 8}}, Uri},
+        definition(Uri, 14, 7)
+    ),
+    %% three: Should go to variable definition (same)
+    ?assertEqual(
+        {#{from => {13, 5}, to => {13, 8}}, Uri},
+        definition(Uri, 15, 5)
+    ),
+    %% four: Should go to first LC generator pattern
+    ?assertEqual(
+        {#{from => {18, 22}, to => {18, 25}}, Uri},
+        definition(Uri, 18, 8)
+    ),
+    %% four: Should go to second LC generator pattern
+    ?assertEqual(
+        {#{from => {19, 22}, to => {19, 26}}, Uri},
+        definition(Uri, 18, 13)
+    ),
+    %% four: Should go to third LC generator pattern
+    ?assertEqual(
+        {#{from => {19, 39}, to => {19, 42}}, Uri},
+        definition(Uri, 19, 32)
     ),
     ok.
 
@@ -716,3 +756,11 @@ parse_incomplete(Config) ->
         els_client:definition(Uri, 19, 3)
     ),
     ok.
+
+definition(Uri, Line, Char) ->
+    case els_client:definition(Uri, Line, Char) of
+        #{result := [#{range := Range, uri := DefUri}]} ->
+            {els_range:to_poi_range(Range), DefUri};
+        Res ->
+            {error, Res}
+    end.
diff --git a/apps/els_lsp/test/els_rename_SUITE.erl b/apps/els_lsp/test/els_rename_SUITE.erl
@@ -18,6 +18,7 @@
     rename_macro/1,
     rename_module/1,
     rename_variable/1,
+    rename_variable_list_comp/1,
     rename_function/1,
     rename_function_quoted_atom/1,
     rename_type/1,
@@ -277,6 +278,99 @@ rename_variable(Config) ->
     assert_changes(Expected10, Result10),
     assert_changes(Expected11, Result11).
 
+-spec rename_variable_list_comp(config()) -> ok.
+rename_variable_list_comp(Config) ->
+    Uri = ?config(variable_list_comp_uri, Config),
+    UriAtom = binary_to_atom(Uri, utf8),
+    NewName = <<"NewAwesomeName">>,
+    %% one: Skip LC
+    Result1 = rename(Uri, 3, 4, NewName),
+    Result1 = rename(Uri, 5, 4, NewName),
+    Expected1 = #{
+        changes => #{
+            UriAtom => [
+                change(NewName, {3, 4}, {3, 7}),
+                change(NewName, {5, 4}, {5, 7})
+            ]
+        }
+    },
+    %% one: Rename in LC only
+    Result2 = rename(Uri, 4, 6, NewName),
+    Result2 = rename(Uri, 4, 13, NewName),
+    Expected2 = #{
+        changes => #{
+            UriAtom => [
+                change(NewName, {4, 6}, {4, 9}),
+                change(NewName, {4, 13}, {4, 16})
+            ]
+        }
+    },
+    %% two: Rename in first LC only
+    Result3 = rename(Uri, 8, 6, NewName),
+    Result3 = rename(Uri, 8, 13, NewName),
+    Expected3 = #{
+        changes => #{
+            UriAtom => [
+                change(NewName, {8, 6}, {8, 9}),
+                change(NewName, {8, 13}, {8, 16})
+            ]
+        }
+    },
+    %% two: Rename in second LC only
+    Result4 = rename(Uri, 9, 6, NewName),
+    Result4 = rename(Uri, 9, 13, NewName),
+    Expected4 = #{
+        changes => #{
+            UriAtom => [
+                change(NewName, {9, 6}, {9, 9}),
+                change(NewName, {9, 13}, {9, 16})
+            ]
+        }
+    },
+    %% three: Rename all Var (no Var in LC pattern)
+    Result5 = rename(Uri, 12, 4, NewName),
+    Result5 = rename(Uri, 13, 6, NewName),
+    Result5 = rename(Uri, 14, 4, NewName),
+    Expected5 = #{
+        changes => #{
+            UriAtom => [
+                change(NewName, {12, 4}, {12, 7}),
+                change(NewName, {13, 6}, {13, 9}),
+                change(NewName, {14, 4}, {14, 7})
+            ]
+        }
+    },
+    %% four: Rename Var2, second LC generator pattern
+    Result6 = rename(Uri, 17, 12, NewName),
+    Result6 = rename(Uri, 18, 21, NewName),
+    Expected6 = #{
+        changes => #{
+            UriAtom => [
+                change(NewName, {17, 12}, {17, 16}),
+                change(NewName, {18, 21}, {18, 25})
+            ]
+        }
+    },
+    %% four: FIXME: Bug, shouldn't rename Var inside second LC
+    %% Result7 = rename(Uri, 17, 7, NewName),
+    %% Result7 = rename(Uri, 17, 21, NewName),
+    %% Expected7 = #{
+    %%     changes => #{
+    %%         UriAtom => [
+    %%             change(NewName, {17, 7}, {17, 10}),
+    %%             change(NewName, {17, 21}, {17, 24})
+    %%         ]
+    %%     }
+    %% },
+    assert_changes(Expected1, Result1),
+    assert_changes(Expected2, Result2),
+    assert_changes(Expected3, Result3),
+    assert_changes(Expected4, Result4),
+    assert_changes(Expected5, Result5),
+    assert_changes(Expected6, Result6),
+    %% assert_changes(Expected7, Result7),
+    ok.
+
 -spec rename_macro(config()) -> ok.
 rename_macro(Config) ->
     Uri = ?config(rename_h_uri, Config),