Improve macro refs performance (#1476)

Improve performance of reference lookup for defines in headers

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

* Use both index and naive when looking up scoped references

Finding scoped references can be done in two ways:

  * Naive, find all POIs that can reach our POI and matches the id.

  * Indexed, use the index to find all matching POIs, then check if
    they actually reference our POI by using goto_definition.

It varies from case to case which is the fastest, so we race both
functions to get the quickest answer.

apps/els_core/src/els_utils.erl

@@ -25,7 +25,8 @@
     camel_case/1,
     jaro_distance/2,
     is_windows/0,
-    system_tmp_dir/0
+    system_tmp_dir/0,
+    race/2
 ]).
 
 %%==============================================================================
@@ -272,9 +273,39 @@ system_tmp_dir() ->
             "/tmp"
     end.
 
+%% @doc Run functions in parallel and return the result of the first function
+%%      that terminates
+-spec race([fun(() -> Result)], timeout()) -> Result.
+race(Funs, Timeout) ->
+    Parent = self(),
+    Ref = make_ref(),
+    Pids = [spawn_link(fun() -> Parent ! {Ref, Fun()} end) || Fun <- Funs],
+    receive
+        {Ref, Result} ->
+            %% Ensure no lingering processes
+            [exit(Pid, kill) || Pid <- Pids],
+            %% Ensure no lingering messages
+            ok = flush(Ref),
+            Result
+    after Timeout ->
+        %% Ensure no lingering processes
+        [exit(Pid, kill) || Pid <- Pids],
+        %% Ensure no lingering messages
+        ok = flush(Ref),
+        error(timeout)
+    end.
+
 %%==============================================================================
 %% Internal functions
 %%==============================================================================
+-spec flush(reference()) -> ok.
+flush(Ref) ->
+    receive
+        {Ref, _} ->
+            flush(Ref)
+    after 0 ->
+        ok
+    end.
 
 %% Folding over files
 

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_field;
 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,12 @@ index_references(Id, Uri, POIs, Version) ->
         %% Type
         type_application,
         %% Atom
-        atom
+        atom,
+        %% Macro
+        macro,
+        %% Record
+        record_expr,
+        record_field
     ],
     [
         index_reference(Id, Uri, POI, Version)
@@ -152,7 +157,10 @@ 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 := {F, A}} = POI, Version) when
+    Kind =/= macro,
+    Kind =/= record_field
+->
     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,79 @@ 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),
-    [
-        {U, Poi}
+-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, POI = #{kind := Kind}) when
+    Kind =:= type_definition
+->
+    %% TODO: This is a hack, ideally we shouldn't have any special handling for
+    %%       these kinds.
+    find_scoped_references_naive(Uri, POI);
+find_scoped_references_for_def(Uri, POI) ->
+    %% Finding scoped references can be done in two ways:
+    %% * Naive, find all POIs that can reach our POI and matches the id.
+    %% * Indexed, use the index to find all matching POIs, then check if
+    %%   they actually reference our POI by using goto_definition.
+    %% It varies from case to case which is the fastest, so we race both
+    %% functions to get the quickest answer.
+    Naive = fun() -> find_scoped_references_naive(Uri, POI) end,
+    Index = fun() -> find_scoped_references_with_index(Uri, POI) end,
+    els_utils:race([Naive, Index], _Timeout = 15000).
+
+-spec find_scoped_references_naive(uri(), els_poi:poi()) -> [location()].
+find_scoped_references_naive(Uri, #{id := Id, kind := Kind}) ->
+    RefKind = kind_to_ref_kind(Kind),
+    Refs = els_scope:local_and_includer_pois(Uri, [RefKind]),
+    MatchingRefs = [
+        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
+    ],
+    ?LOG_DEBUG(
+        "Found scoped references (naive) for ~p: ~p",
+        [Id, length(MatchingRefs)]
+    ),
+    MatchingRefs.
+
+-spec find_scoped_references_with_index(uri(), els_poi:poi()) -> [location()].
+find_scoped_references_with_index(Uri, POI = #{kind := Kind, id := Id}) ->
+    RefPOI = POI#{kind := kind_to_ref_kind(Kind)},
+    Match = fun(#{uri := RefUri}) ->
+        case els_code_navigation:goto_definition(RefUri, RefPOI) of
+            {ok, [{Uri, _}]} -> true;
+            _Else -> false
+        end
+    end,
+    Refs = [Ref || Ref <- find_references_for_id(Kind, Id), Match(Ref)],
+    ?LOG_DEBUG(
+        "Found scoped references (with index) for ~p: ~p",
+        [Id, length(Refs)]
+    ),
+    Refs.
+
+-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 +248,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,10 @@ extract_pattern({include_lib, Id}) ->
 extract_pattern({behaviour, Name}) ->
     Name;
 extract_pattern({atom, Name}) ->
+    Name;
+extract_pattern({record_field, {_Record, Name}}) ->
+    Name;
+extract_pattern({record, Name}) ->
     Name.
 
 -spec include_id(string()) -> string().

elvis.config

@@ -21,7 +21,9 @@
                             els_hover_SUITE,
                             els_parser_SUITE,
                             els_references_SUITE,
-                            els_methods
+                            els_methods,
+                            %% TODO: We should probably split this up
+                            els_utils
                         ]
                     }},
                     {elvis_style, dont_repeat_yourself, #{