Improve performance of reference lookup for defines in headers

Use index to find references instead of traversing all files that
includes a header.

apps/els_lsp/src/els_dt_document.erl

@@ -283,24 +283,28 @@ find_candidates(Pattern) ->
 get_words(Text) ->
     case erl_scan:string(els_utils:to_list(Text)) of
         {ok, Tokens, _EndLocation} ->
-            Fun = fun
-                ({atom, _Location, Atom}, Words) ->
-                    sets:add_element(Atom, Words);
-                ({string, _Location, String}, Words) ->
-                    case filename:extension(String) of
-                        ".hrl" ->
-                            Id = filename:rootname(filename:basename(String)),
-                            sets:add_element(Id, Words);
-                        _ ->
-                            Words
-                    end;
-                (_, Words) ->
-                    Words
-            end,
-            lists:foldl(Fun, sets:new(), Tokens);
+            tokens_to_words(Tokens, sets:new());
         {error, ErrorInfo, ErrorLocation} ->
             ?LOG_WARNING("Errors while get_words [info=~p] [location=~p]", [
                 ErrorInfo, ErrorLocation
             ]),
             sets:new()
     end.
+
+-spec tokens_to_words([erl_scan:token()], sets:set()) -> sets:set().
+tokens_to_words([{atom, _Location, Atom} | Tokens], Words) ->
+    tokens_to_words(Tokens, sets:add_element(Atom, Words));
+tokens_to_words([{string, _Location, String} | Tokens], Words) ->
+    case filename:extension(String) of
+        ".hrl" ->
+            Id = filename:rootname(filename:basename(String)),
+            tokens_to_words(Tokens, sets:add_element(Id, Words));
+        _ ->
+            tokens_to_words(Tokens, Words)
+    end;
+tokens_to_words([{'?', _}, {var, _, Macro} | Tokens], Words) ->
+    tokens_to_words(Tokens, sets:add_element(Macro, Words));
+tokens_to_words([_ | Tokens], Words) ->
+    tokens_to_words(Tokens, Words);
+tokens_to_words([], Words) ->
+    Words.

apps/els_lsp/src/els_dt_references.erl

@@ -181,6 +181,11 @@ kind_to_category(Kind) when
     Kind =:= record
 ->
     record;
+kind_to_category(Kind) when
+    Kind =:= record_field;
+    Kind =:= record_def_field
+->
+    record;
 kind_to_category(Kind) when Kind =:= include ->
     include;
 kind_to_category(Kind) when Kind =:= include_lib ->

apps/els_lsp/src/els_indexing.erl

@@ -142,7 +142,11 @@ index_references(Id, Uri, POIs, Version) ->
         %% Type
         type_application,
         %% Atom
-        atom
+        atom,
+        %% Macro
+        macro,
+        %% Record
+        record_expr
     ],
     [
         index_reference(Id, Uri, POI, Version)
@@ -152,7 +156,16 @@ index_references(Id, Uri, POIs, Version) ->
     ok.
 
 -spec index_reference(atom(), uri(), els_poi:poi(), version()) -> ok.
-index_reference(M, Uri, #{id := {F, A}} = POI, Version) ->
+index_reference(_M, Uri, #{kind := Kind, id := Id, range := Range}, Version) when
+    Kind =:= macro
+->
+    els_dt_references:versioned_insert(Kind, #{
+        id => Id,
+        uri => Uri,
+        range => Range,
+        version => Version
+    });
+index_reference(M, Uri, #{kind := _Kind, id := {F, A}} = POI, Version) ->
     index_reference(M, Uri, POI#{id => {M, F, A}}, Version);
 index_reference(_M, Uri, #{kind := Kind, id := Id, range := Range}, Version) ->
     els_dt_references:versioned_insert(Kind, #{

apps/els_lsp/src/els_references_provider.erl

@@ -101,14 +101,12 @@ find_references(Uri, #{
     {F, A, _Index} = Id,
     Key = {els_uri:module(Uri), F, A},
     find_references_for_id(Kind, Key);
-find_references(Uri, Poi = #{kind := Kind}) when
+find_references(Uri, POI = #{kind := Kind}) when
     Kind =:= record;
     Kind =:= record_def_field;
     Kind =:= define
 ->
-    uri_pois_to_locations(
-        find_scoped_references_for_def(Uri, Poi)
-    );
+    find_scoped_references_for_def(Uri, POI);
 find_references(Uri, Poi = #{kind := Kind, id := Id}) when
     Kind =:= type_definition
 ->
@@ -119,11 +117,9 @@ find_references(Uri, Poi = #{kind := Kind, id := Id}) when
         end,
     lists:usort(
         find_references_for_id(Kind, Key) ++
-            uri_pois_to_locations(
-                find_scoped_references_for_def(Uri, Poi)
-            )
+            find_scoped_references_for_def(Uri, Poi)
     );
-find_references(Uri, Poi = #{kind := Kind}) when
+find_references(Uri, Poi = #{kind := Kind, id := Id}) when
     Kind =:= record_expr;
     Kind =:= record_field;
     Kind =:= macro;
@@ -134,9 +130,7 @@ find_references(Uri, Poi = #{kind := Kind}) when
             find_references(DefUri, DefPoi);
         _ ->
             %% look for references only in the current document
-            uri_pois_to_locations(
-                find_scoped_references_for_def(Uri, Poi)
-            )
+            local_refs(Uri, Kind, Id)
     end;
 find_references(Uri, #{kind := module}) ->
     Refs = find_references_to_module(Uri),
@@ -149,28 +143,54 @@ find_references(_Uri, #{kind := Kind, id := Name}) when
 find_references(_Uri, _POI) ->
     [].
 
--spec find_scoped_references_for_def(uri(), els_poi:poi()) -> [{uri(), els_poi:poi()}].
-find_scoped_references_for_def(Uri, #{kind := Kind, id := Name}) ->
-    Kinds = kind_to_ref_kinds(Kind),
-    Refs = els_scope:local_and_includer_pois(Uri, Kinds),
+-spec local_refs(uri(), els_poi:poi_kind(), els_poi:poi_id()) ->
+    [location()].
+local_refs(Uri, Kind, Id) ->
+    {ok, Document} = els_utils:lookup_document(Uri),
+    POIs = els_dt_document:pois(Document, [kind_to_ref_kind(Kind)]),
+    LocalRefs = [
+        location(Uri, R)
+     || #{range := R, id := IdPoi} <- POIs,
+        Id == IdPoi
+    ],
+    LocalRefs.
+
+-spec find_scoped_references_for_def(uri(), els_poi:poi()) -> [location()].
+find_scoped_references_for_def(Uri, #{kind := Kind, id := Id}) when
+    Kind =:= record_def_field;
+    Kind =:= type_definition
+->
+    %% TODO: This is a hack, ideally we shouldn't have any special handling for
+    %%       these kinds
+    RefKind = kind_to_ref_kind(Kind),
+    Refs = els_scope:local_and_includer_pois(Uri, [RefKind]),
     [
-        {U, Poi}
+        location(U, R)
      || {U, Pois} <- Refs,
-        #{id := N} = Poi <- Pois,
-        N =:= Name
-    ].
-
--spec kind_to_ref_kinds(els_poi:poi_kind()) -> [els_poi:poi_kind()].
-kind_to_ref_kinds(define) ->
-    [macro];
-kind_to_ref_kinds(record) ->
-    [record_expr];
-kind_to_ref_kinds(record_def_field) ->
-    [record_field];
-kind_to_ref_kinds(type_definition) ->
-    [type_application];
-kind_to_ref_kinds(Kind) ->
-    [Kind].
+        #{id := N, range := R} <- Pois,
+        N =:= Id
+    ];
+find_scoped_references_for_def(Uri, POI = #{kind := Kind, id := Id}) ->
+    RefPOI = POI#{kind := kind_to_ref_kind(Kind)},
+    F = fun(#{uri := RefUri}) ->
+        case els_code_navigation:goto_definition(RefUri, RefPOI) of
+            {ok, Uri, _} -> true;
+            _ -> false
+        end
+    end,
+    [Ref || Ref <- find_references_for_id(Kind, Id), F(Ref)].
+
+-spec kind_to_ref_kind(els_poi:poi_kind()) -> els_poi:poi_kind().
+kind_to_ref_kind(define) ->
+    macro;
+kind_to_ref_kind(record) ->
+    record_expr;
+kind_to_ref_kind(record_def_field) ->
+    record_field;
+kind_to_ref_kind(type_definition) ->
+    type_application;
+kind_to_ref_kind(Kind) ->
+    Kind.
 
 -spec find_references_to_module(uri()) -> [els_dt_references:item()].
 find_references_to_module(Uri) ->
@@ -203,10 +223,6 @@ find_references_for_id(Kind, Id) ->
     {ok, Refs} = els_dt_references:find_by_id(Kind, Id),
     [location(U, R) || #{uri := U, range := R} <- Refs].
 
--spec uri_pois_to_locations([{uri(), els_poi:poi()}]) -> [location()].
-uri_pois_to_locations(Refs) ->
-    [location(U, R) || {U, #{range := R}} <- Refs].
-
 -spec location(uri(), els_poi:poi_range()) -> location().
 location(Uri, Range) ->
     #{

apps/els_lsp/src/els_rename_provider.erl

@@ -307,10 +307,10 @@ changes(Uri, #{kind := DefKind} = DefPoi, NewName) when
     Self = #{range => editable_range(DefPoi), newText => NewName},
     Refs = els_references_provider:find_scoped_references_for_def(Uri, DefPoi),
     lists:foldl(
-        fun({U, Poi}, Acc) ->
+        fun(#{uri := U, range := R}, Acc) ->
             Change = #{
-                range => editable_range(Poi),
-                newText => new_name(Poi, NewName)
+                range => R,
+                newText => new_name(DefKind, NewName)
             },
             maps:update_with(U, fun(V) -> [Change | V] end, [Change], Acc)
         end,
@@ -320,10 +320,10 @@ changes(Uri, #{kind := DefKind} = DefPoi, NewName) when
 changes(_Uri, _POI, _NewName) ->
     null.
 
--spec new_name(els_poi:poi(), binary()) -> binary().
-new_name(#{kind := macro}, NewName) ->
+-spec new_name(els_poi:poi_kind(), binary()) -> binary().
+new_name(define, NewName) ->
     <<"?", NewName/binary>>;
-new_name(#{kind := record_expr}, NewName) ->
+new_name(record, NewName) ->
     <<"#", NewName/binary>>;
 new_name(_, NewName) ->
     NewName.

apps/els_lsp/src/els_text_search.erl

@@ -28,6 +28,8 @@ find_candidate_uris(Id) ->
     atom() | binary().
 extract_pattern({function, {_M, F, _A}}) ->
     F;
+extract_pattern({type, {F, _A}}) ->
+    F;
 extract_pattern({type, {_M, F, _A}}) ->
     F;
 extract_pattern({macro, {Name, _Arity}}) ->
@@ -41,6 +43,11 @@ extract_pattern({include_lib, Id}) ->
 extract_pattern({behaviour, Name}) ->
     Name;
 extract_pattern({atom, Name}) ->
+    Name;
+extract_pattern({record, {_Record, Name}}) ->
+    %% Record fields
+    Name;
+extract_pattern({record, Name}) ->
     Name.
 
 -spec include_id(string()) -> string().