Ensure uniqueness of file ids in reference search via hashmap

9 files changed, 122 insertions, 121 deletions

diff --git a/crates/assists/src/handlers/extract_struct_from_enum_variant.rs b/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
index 21b13977b..e3ef04932 100644
--- a/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
+++ b/crates/assists/src/handlers/extract_struct_from_enum_variant.rs
@@ -8,7 +8,7 @@ use ide_db::{
         insert_use::{insert_use, ImportScope},
         mod_path_to_ast,
     },
-    search::{FileReference, FileReferences},
+    search::FileReference,
     RootDatabase,
 };
 use rustc_hash::FxHashSet;
@@ -63,10 +63,10 @@ pub(crate) fn extract_struct_from_enum_variant(
             let current_module = enum_hir.module(ctx.db());
             visited_modules_set.insert(current_module);
             let mut def_rewriter = None;
-            for FileReferences { file_id, references: refs } in usages {
+            for (file_id, references) in usages {
                 let mut rewriter = SyntaxRewriter::default();
                 let source_file = ctx.sema.parse(file_id);
-                for reference in refs {
+                for reference in references {
                     update_reference(
                         ctx,
                         &mut rewriter,
diff --git a/crates/assists/src/handlers/inline_local_variable.rs b/crates/assists/src/handlers/inline_local_variable.rs
index 928df6825..dc798daaa 100644
--- a/crates/assists/src/handlers/inline_local_variable.rs
+++ b/crates/assists/src/handlers/inline_local_variable.rs
@@ -47,8 +47,8 @@ pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> O
 
     let def = ctx.sema.to_def(&bind_pat)?;
     let def = Definition::Local(def);
-    let refs = def.usages(&ctx.sema).all();
-    if refs.is_empty() {
+    let usages = def.usages(&ctx.sema).all();
+    if usages.is_empty() {
         mark::hit!(test_not_applicable_if_variable_unused);
         return None;
     };
@@ -66,9 +66,10 @@ pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> O
         let_stmt.syntax().text_range()
     };
 
-    let wrap_in_parens = refs
-        .iter()
-        .flat_map(|refs| &refs.references)
+    let wrap_in_parens = usages
+        .references
+        .values()
+        .flatten()
         .map(|&FileReference { range, .. }| {
             let usage_node =
                 ctx.covering_node_for_range(range).ancestors().find_map(ast::PathExpr::cast)?;
@@ -115,8 +116,7 @@ pub(crate) fn inline_local_variable(acc: &mut Assists, ctx: &AssistContext) -> O
         target,
         move |builder| {
             builder.delete(delete_range);
-            for (reference, should_wrap) in
-                refs.iter().flat_map(|refs| &refs.references).zip(wrap_in_parens)
+            for (reference, should_wrap) in usages.references.values().flatten().zip(wrap_in_parens)
             {
                 let replacement =
                     if should_wrap { init_in_paren.clone() } else { init_str.clone() };
diff --git a/crates/assists/src/handlers/remove_unused_param.rs b/crates/assists/src/handlers/remove_unused_param.rs
index 4f3b8ac46..c961680e2 100644
--- a/crates/assists/src/handlers/remove_unused_param.rs
+++ b/crates/assists/src/handlers/remove_unused_param.rs
@@ -1,7 +1,4 @@
-use ide_db::{
-    defs::Definition,
-    search::{FileReference, FileReferences},
-};
+use ide_db::{base_db::FileId, defs::Definition, search::FileReference};
 use syntax::{
     algo::find_node_at_range,
     ast::{self, ArgListOwner},
@@ -61,8 +58,8 @@ pub(crate) fn remove_unused_param(acc: &mut Assists, ctx: &AssistContext) -> Opt
         param.syntax().text_range(),
         |builder| {
             builder.delete(range_to_remove(param.syntax()));
-            for usages in fn_def.usages(&ctx.sema).all() {
-                process_usages(ctx, builder, usages, param_position);
+            for (file_id, references) in fn_def.usages(&ctx.sema).all() {
+                process_usages(ctx, builder, file_id, references, param_position);
             }
         },
     )
@@ -71,12 +68,13 @@ pub(crate) fn remove_unused_param(acc: &mut Assists, ctx: &AssistContext) -> Opt
 fn process_usages(
     ctx: &AssistContext,
     builder: &mut AssistBuilder,
-    usages: FileReferences,
+    file_id: FileId,
+    references: Vec<FileReference>,
     arg_to_remove: usize,
 ) {
-    let source_file = ctx.sema.parse(usages.file_id);
-    builder.edit_file(usages.file_id);
-    for usage in usages.references {
+    let source_file = ctx.sema.parse(file_id);
+    builder.edit_file(file_id);
+    for usage in references {
         if let Some(text_range) = process_usage(&source_file, usage, arg_to_remove) {
             builder.delete(text_range);
         }
diff --git a/crates/ide/src/call_hierarchy.rs b/crates/ide/src/call_hierarchy.rs
index 90d3b9a31..e8999a7f3 100644
--- a/crates/ide/src/call_hierarchy.rs
+++ b/crates/ide/src/call_hierarchy.rs
@@ -3,8 +3,8 @@
 use indexmap::IndexMap;
 
 use hir::Semantics;
+use ide_db::call_info::FnCallNode;
 use ide_db::RootDatabase;
-use ide_db::{call_info::FnCallNode, search::FileReferences};
 use syntax::{ast, AstNode, TextRange};
 
 use crate::{
@@ -47,7 +47,7 @@ pub(crate) fn incoming_calls(db: &RootDatabase, position: FilePosition) -> Optio
 
     let mut calls = CallLocations::default();
 
-    for &FileReferences { file_id, ref references } in refs.info.references() {
+    for (&file_id, references) in refs.info.references().iter() {
         let file = sema.parse(file_id);
         let file = file.syntax();
         for reference in references {
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index 132680bfb..c7943dc95 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -13,6 +13,7 @@ pub(crate) mod rename;
 
 use hir::Semantics;
 use ide_db::{
+    base_db::FileId,
     defs::{Definition, NameClass, NameRefClass},
     search::{FileReference, FileReferences, ReferenceAccess, ReferenceKind, SearchScope},
     RootDatabase,
@@ -28,7 +29,7 @@ use crate::{display::TryToNav, FilePosition, FileRange, NavigationTarget, RangeI
 #[derive(Debug, Clone)]
 pub struct ReferenceSearchResult {
     declaration: Declaration,
-    references: Vec<FileReferences>,
+    references: FileReferences,
 }
 
 #[derive(Debug, Clone)]
@@ -47,10 +48,21 @@ impl ReferenceSearchResult {
         &self.declaration.nav
     }
 
-    pub fn references(&self) -> &[FileReferences] {
+    pub fn references(&self) -> &FileReferences {
         &self.references
     }
 
+    pub fn references_with_declaration(mut self) -> FileReferences {
+        let decl_ref = FileReference {
+            range: self.declaration.nav.focus_or_full_range(),
+            kind: self.declaration.kind,
+            access: self.declaration.access,
+        };
+        let file_id = self.declaration.nav.file_id;
+        self.references.references.entry(file_id).or_default().push(decl_ref);
+        self.references
+    }
+
     /// Total number of references
     /// At least 1 since all valid references should
     /// Have a declaration
@@ -62,23 +74,11 @@ impl ReferenceSearchResult {
 // allow turning ReferenceSearchResult into an iterator
 // over References
 impl IntoIterator for ReferenceSearchResult {
-    type Item = FileReferences;
-    type IntoIter = std::vec::IntoIter<FileReferences>;
+    type Item = (FileId, Vec<FileReference>);
+    type IntoIter = std::collections::hash_map::IntoIter<FileId, Vec<FileReference>>;
 
-    fn into_iter(mut self) -> Self::IntoIter {
-        let mut v = Vec::with_capacity(self.len());
-        v.append(&mut self.references);
-        let decl_ref = FileReference {
-            range: self.declaration.nav.focus_or_full_range(),
-            kind: self.declaration.kind,
-            access: self.declaration.access,
-        };
-        let file_id = self.declaration.nav.file_id;
-        match v.iter_mut().find(|it| it.file_id == file_id) {
-            Some(file_refs) => file_refs.references.push(decl_ref),
-            None => v.push(FileReferences { file_id, references: vec![decl_ref] }),
-        }
-        v.into_iter()
+    fn into_iter(self) -> Self::IntoIter {
+        self.references_with_declaration().into_iter()
     }
 }
 
@@ -110,11 +110,12 @@ pub(crate) fn find_all_refs(
 
     let RangeInfo { range, info: def } = find_name(&sema, &syntax, position, opt_name)?;
 
-    let mut references = def.usages(sema).set_scope(search_scope).all();
-    references.iter_mut().for_each(|it| {
-        it.references.retain(|r| search_kind == ReferenceKind::Other || search_kind == r.kind)
-    });
-    references.retain(|r| !r.references.is_empty());
+    let mut usages = def.usages(sema).set_scope(search_scope).all();
+    usages
+        .references
+        .values_mut()
+        .for_each(|it| it.retain(|r| search_kind == ReferenceKind::Other || search_kind == r.kind));
+    usages.references.retain(|_, it| !it.is_empty());
 
     let nav = def.try_to_nav(sema.db)?;
     let decl_range = nav.focus_or_full_range();
@@ -138,7 +139,7 @@ pub(crate) fn find_all_refs(
 
     let declaration = Declaration { nav, kind, access: decl_access(&def, &syntax, decl_range) };
 
-    Some(RangeInfo::new(range, ReferenceSearchResult { declaration, references }))
+    Some(RangeInfo::new(range, ReferenceSearchResult { declaration, references: usages }))
 }
 
 fn find_name(
@@ -292,32 +293,30 @@ fn try_find_self_references(
             ReferenceAccess::Read
         }),
     };
-    let references = function
+    let refs = function
         .body()
         .map(|body| {
-            FileReferences {
-                file_id,
-                references: body
-                    .syntax()
-                    .descendants()
-                    .filter_map(ast::PathExpr::cast)
-                    .filter_map(|expr| {
-                        let path = expr.path()?;
-                        if path.qualifier().is_none() {
-                            path.segment()?.self_token()
-                        } else {
-                            None
-                        }
-                    })
-                    .map(|token| FileReference {
-                        range: token.text_range(),
-                        kind: ReferenceKind::SelfKw,
-                        access: declaration.access, // FIXME: properly check access kind here instead of copying it from the declaration
-                    })
-                    .collect(),
-            }
+            body.syntax()
+                .descendants()
+                .filter_map(ast::PathExpr::cast)
+                .filter_map(|expr| {
+                    let path = expr.path()?;
+                    if path.qualifier().is_none() {
+                        path.segment()?.self_token()
+                    } else {
+                        None
+                    }
+                })
+                .map(|token| FileReference {
+                    range: token.text_range(),
+                    kind: ReferenceKind::SelfKw,
+                    access: declaration.access, // FIXME: properly check access kind here instead of copying it from the declaration
+                })
+                .collect()
         })
-        .map_or_else(Vec::default, |it| vec![it]);
+        .unwrap_or_default();
+    let mut references = FileReferences::default();
+    references.references.insert(file_id, refs);
 
     Some(RangeInfo::new(
         param_self_token.text_range(),
@@ -328,7 +327,7 @@ fn try_find_self_references(
 #[cfg(test)]
 mod tests {
     use expect_test::{expect, Expect};
-    use ide_db::{base_db::FileId, search::FileReferences};
+    use ide_db::base_db::FileId;
     use stdx::format_to;
 
     use crate::{fixture, SearchScope};
@@ -1022,7 +1021,7 @@ impl Foo {
             actual += "\n\n";
         }
 
-        for FileReferences { file_id, references } in refs.references {
+        for (file_id, references) in refs.references {
             for r in references {
                 format_to!(actual, "{:?} {:?} {:?}", file_id, r.range, r.kind);
                 if let Some(access) = r.access {
diff --git a/crates/ide/src/references/rename.rs b/crates/ide/src/references/rename.rs
index dd08e1c32..5207388b5 100644
--- a/crates/ide/src/references/rename.rs
+++ b/crates/ide/src/references/rename.rs
@@ -9,7 +9,7 @@ use hir::{Module, ModuleDef, ModuleSource, Semantics};
 use ide_db::{
     base_db::{AnchoredPathBuf, FileId, FileRange, SourceDatabaseExt},
     defs::{Definition, NameClass, NameRefClass},
-    search::FileReferences,
+    search::FileReference,
     RootDatabase,
 };
 use syntax::{
@@ -176,7 +176,8 @@ fn find_all_refs(
 
 fn source_edit_from_references(
     sema: &Semantics<RootDatabase>,
-    &FileReferences { file_id, ref references }: &FileReferences,
+    file_id: FileId,
+    references: &[FileReference],
     new_name: &str,
 ) -> SourceFileEdit {
     let mut edit = TextEdit::builder();
@@ -283,10 +284,9 @@ fn rename_mod(
     }
 
     let RangeInfo { range, info: refs } = find_all_refs(sema, position)?;
-    let ref_edits = refs
-        .references()
-        .iter()
-        .map(|reference| source_edit_from_references(sema, reference, new_name));
+    let ref_edits = refs.references().iter().map(|(&file_id, references)| {
+        source_edit_from_references(sema, file_id, references, new_name)
+    });
     source_file_edits.extend(ref_edits);
 
     Ok(RangeInfo::new(range, SourceChange::from_edits(source_file_edits, file_system_edits)))
@@ -341,7 +341,9 @@ fn rename_to_self(
     let mut edits = refs
         .references()
         .iter()
-        .map(|reference| source_edit_from_references(sema, reference, "self"))
+        .map(|(&file_id, references)| {
+            source_edit_from_references(sema, file_id, references, "self")
+        })
         .collect::<Vec<_>>();
 
     edits.push(SourceFileEdit {
@@ -467,7 +469,9 @@ fn rename_reference(
 
     let edit = refs
         .into_iter()
-        .map(|reference| source_edit_from_references(sema, &reference, new_name))
+        .map(|(file_id, references)| {
+            source_edit_from_references(sema, file_id, &references, new_name)
+        })
         .collect::<Vec<_>>();
 
     Ok(RangeInfo::new(range, SourceChange::from(edit)))
diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index b8359a9b4..89a313e9b 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -8,7 +8,6 @@ use std::{convert::TryInto, mem};
 
 use base_db::{FileId, FileRange, SourceDatabaseExt};
 use hir::{DefWithBody, HasSource, Module, ModuleSource, Semantics, Visibility};
-use itertools::Itertools;
 use once_cell::unsync::Lazy;
 use rustc_hash::FxHashMap;
 use syntax::{ast, match_ast, AstNode, TextRange, TextSize};
@@ -19,17 +18,37 @@ use crate::{
     RootDatabase,
 };
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Default, Clone)]
 pub struct FileReferences {
-    pub file_id: FileId,
-    pub references: Vec<FileReference>,
+    pub references: FxHashMap<FileId, Vec<FileReference>>,
 }
 
 impl FileReferences {
+    pub fn is_empty(&self) -> bool {
+        self.references.is_empty()
+    }
+
+    pub fn len(&self) -> usize {
+        self.references.len()
+    }
+
+    pub fn iter(&self) -> impl Iterator<Item = (&FileId, &Vec<FileReference>)> + '_ {
+        self.references.iter()
+    }
+
     pub fn file_ranges(&self) -> impl Iterator<Item = FileRange> + '_ {
-        self.references
-            .iter()
-            .map(move |&FileReference { range, .. }| FileRange { file_id: self.file_id, range })
+        self.references.iter().flat_map(|(&file_id, refs)| {
+            refs.iter().map(move |&FileReference { range, .. }| FileRange { file_id, range })
+        })
+    }
+}
+
+impl IntoIterator for FileReferences {
+    type Item = (FileId, Vec<FileReference>);
+    type IntoIter = <FxHashMap<FileId, Vec<FileReference>> as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.references.into_iter()
     }
 }
 
@@ -275,21 +294,12 @@ impl<'a> FindUsages<'a> {
     }
 
     /// The [`FileReferences`] returned always have unique [`FileId`]s.
-    pub fn all(self) -> Vec<FileReferences> {
-        let mut res = <Vec<FileReferences>>::new();
+    pub fn all(self) -> FileReferences {
+        let mut res = FileReferences::default();
         self.search(&mut |file_id, reference| {
-            match res.iter_mut().find(|it| it.file_id == file_id) {
-                Some(file_refs) => file_refs.references.push(reference),
-                _ => res.push(FileReferences { file_id, references: vec![reference] }),
-            }
+            res.references.entry(file_id).or_default().push(reference);
             false
         });
-        assert!(res
-            .iter()
-            .map(|refs| refs.file_id)
-            .sorted_unstable()
-            .tuple_windows::<(_, _)>()
-            .all(|(a, b)| a < b));
         res
     }
 
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index d862f370a..2cc57f022 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -12,7 +12,6 @@ use ide::{
     FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData, LineIndex, NavigationTarget,
     Query, RangeInfo, Runnable, RunnableKind, SearchScope, SourceChange, SymbolKind, TextEdit,
 };
-use ide_db::search::FileReferences;
 use itertools::Itertools;
 use lsp_server::ErrorCode;
 use lsp_types::{
@@ -813,18 +812,14 @@ pub(crate) fn handle_references(
     };
 
     let locations = if params.context.include_declaration {
-        let mut locations = Vec::default();
-        refs.into_iter().for_each(|it| {
-            locations.extend(
-                it.file_ranges().filter_map(|frange| to_proto::location(&snap, frange).ok()),
-            )
-        });
-        locations
+        refs.references_with_declaration()
+            .file_ranges()
+            .filter_map(|frange| to_proto::location(&snap, frange).ok())
+            .collect()
     } else {
         // Only iterate over the references if include_declaration was false
         refs.references()
-            .iter()
-            .flat_map(FileReferences::file_ranges)
+            .file_ranges()
             .filter_map(|frange| to_proto::location(&snap, frange).ok())
             .collect()
     };
@@ -1181,8 +1176,7 @@ pub(crate) fn handle_code_lens_resolve(
                 .unwrap_or(None)
                 .map(|r| {
                     r.references()
-                        .iter()
-                        .flat_map(FileReferences::file_ranges)
+                        .file_ranges()
                         .filter_map(|frange| to_proto::location(&snap, frange).ok())
                         .collect_vec()
                 })
@@ -1227,11 +1221,11 @@ pub(crate) fn handle_document_highlight(
     };
 
     let res = refs
-        .into_iter()
-        .find(|refs| refs.file_id == position.file_id)
+        .references_with_declaration()
+        .references
+        .get(&position.file_id)
         .map(|file_refs| {
             file_refs
-                .references
                 .into_iter()
                 .map(|r| DocumentHighlight {
                     range: to_proto::range(&line_index, r.range),
diff --git a/crates/ssr/src/search.rs b/crates/ssr/src/search.rs
index a1d653aff..a3eb2e800 100644
--- a/crates/ssr/src/search.rs
+++ b/crates/ssr/src/search.rs
@@ -20,7 +20,7 @@ use test_utils::mark;
 /// them more than once.
 #[derive(Default)]
 pub(crate) struct UsageCache {
-    usages: Vec<(Definition, Vec<FileReferences>)>,
+    usages: Vec<(Definition, FileReferences)>,
 }
 
 impl<'db> MatchFinder<'db> {
@@ -58,11 +58,7 @@ impl<'db> MatchFinder<'db> {
     ) {
         if let Some(resolved_path) = pick_path_for_usages(pattern) {
             let definition: Definition = resolved_path.resolution.clone().into();
-            for file_range in self
-                .find_usages(usage_cache, definition)
-                .iter()
-                .flat_map(FileReferences::file_ranges)
-            {
+            for file_range in self.find_usages(usage_cache, definition).file_ranges() {
                 if let Some(node_to_match) = self.find_node_to_match(resolved_path, file_range) {
                     if !is_search_permitted_ancestors(&node_to_match) {
                         mark::hit!(use_declaration_with_braces);
@@ -112,7 +108,7 @@ impl<'db> MatchFinder<'db> {
         &self,
         usage_cache: &'a mut UsageCache,
         definition: Definition,
-    ) -> &'a [FileReferences] {
+    ) -> &'a FileReferences {
         // Logically if a lookup succeeds we should just return it. Unfortunately returning it would
         // extend the lifetime of the borrow, then we wouldn't be able to do the insertion on a
         // cache miss. This is a limitation of NLL and is fixed with Polonius. For now we do two
@@ -254,7 +250,7 @@ fn is_search_permitted(node: &SyntaxNode) -> bool {
 }
 
 impl UsageCache {
-    fn find(&mut self, definition: &Definition) -> Option<&[FileReferences]> {
+    fn find(&mut self, definition: &Definition) -> Option<&FileReferences> {
         // We expect a very small number of cache entries (generally 1), so a linear scan should be
         // fast enough and avoids the need to implement Hash for Definition.
         for (d, refs) in &self.usages {