4 files changed, 110 insertions, 100 deletions

diff --git a/crates/ide/src/call_hierarchy.rs b/crates/ide/src/call_hierarchy.rs
index b29d1fef9..e8999a7f3 100644
--- a/crates/ide/src/call_hierarchy.rs
+++ b/crates/ide/src/call_hierarchy.rs
@@ -47,22 +47,23 @@ pub(crate) fn incoming_calls(db: &RootDatabase, position: FilePosition) -> Optio
 
     let mut calls = CallLocations::default();
 
-    for reference in refs.info.references() {
-        let file_id = reference.file_range.file_id;
+    for (&file_id, references) in refs.info.references().iter() {
         let file = sema.parse(file_id);
         let file = file.syntax();
-        let token = file.token_at_offset(reference.file_range.range.start()).next()?;
-        let token = sema.descend_into_macros(token);
-        let syntax = token.parent();
-
-        // This target is the containing function
-        if let Some(nav) = syntax.ancestors().find_map(|node| {
-            let fn_ = ast::Fn::cast(node)?;
-            let def = sema.to_def(&fn_)?;
-            def.try_to_nav(sema.db)
-        }) {
-            let relative_range = reference.file_range.range;
-            calls.add(&nav, relative_range);
+        for reference in references {
+            let token = file.token_at_offset(reference.range.start()).next()?;
+            let token = sema.descend_into_macros(token);
+            let syntax = token.parent();
+
+            // This target is the containing function
+            if let Some(nav) = syntax.ancestors().find_map(|node| {
+                let fn_ = ast::Fn::cast(node)?;
+                let def = sema.to_def(&fn_)?;
+                def.try_to_nav(sema.db)
+            }) {
+                let relative_range = reference.range;
+                calls.add(&nav, relative_range);
+            }
         }
     }
 
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 1f368cbd0..1e03832ec 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -92,7 +92,7 @@ pub use ide_db::base_db::{
 };
 pub use ide_db::{
     call_info::CallInfo,
-    search::{Reference, ReferenceAccess, ReferenceKind},
+    search::{FileReference, ReferenceAccess, ReferenceKind},
 };
 pub use ide_db::{
     label::Label,
diff --git a/crates/ide/src/references.rs b/crates/ide/src/references.rs
index b774a2be1..7d4757e02 100644
--- a/crates/ide/src/references.rs
+++ b/crates/ide/src/references.rs
@@ -13,9 +13,9 @@ pub(crate) mod rename;
 
 use hir::Semantics;
 use ide_db::{
+    base_db::FileId,
     defs::{Definition, NameClass, NameRefClass},
-    search::Reference,
-    search::{ReferenceAccess, ReferenceKind, SearchScope},
+    search::{FileReference, ReferenceAccess, ReferenceKind, SearchScope, UsageSearchResult},
     RootDatabase,
 };
 use syntax::{
@@ -29,7 +29,7 @@ use crate::{display::TryToNav, FilePosition, FileRange, NavigationTarget, RangeI
 #[derive(Debug, Clone)]
 pub struct ReferenceSearchResult {
     declaration: Declaration,
-    references: Vec<Reference>,
+    references: UsageSearchResult,
 }
 
 #[derive(Debug, Clone)]
@@ -48,10 +48,21 @@ impl ReferenceSearchResult {
         &self.declaration.nav
     }
 
-    pub fn references(&self) -> &[Reference] {
+    pub fn references(&self) -> &UsageSearchResult {
         &self.references
     }
 
+    pub fn references_with_declaration(mut self) -> UsageSearchResult {
+        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
@@ -63,21 +74,11 @@ impl ReferenceSearchResult {
 // allow turning ReferenceSearchResult into an iterator
 // over References
 impl IntoIterator for ReferenceSearchResult {
-    type Item = Reference;
-    type IntoIter = std::vec::IntoIter<Reference>;
-
-    fn into_iter(mut self) -> Self::IntoIter {
-        let mut v = Vec::with_capacity(self.len());
-        v.push(Reference {
-            file_range: FileRange {
-                file_id: self.declaration.nav.file_id,
-                range: self.declaration.nav.focus_or_full_range(),
-            },
-            kind: self.declaration.kind,
-            access: self.declaration.access,
-        });
-        v.append(&mut self.references);
-        v.into_iter()
+    type Item = (FileId, Vec<FileReference>);
+    type IntoIter = std::collections::hash_map::IntoIter<FileId, Vec<FileReference>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.references_with_declaration().into_iter()
     }
 }
 
@@ -109,13 +110,12 @@ pub(crate) fn find_all_refs(
 
     let RangeInfo { range, info: def } = find_name(&sema, &syntax, position, opt_name)?;
 
-    let references = def
-        .usages(sema)
-        .set_scope(search_scope)
-        .all()
-        .into_iter()
-        .filter(|r| search_kind == ReferenceKind::Other || search_kind == r.kind)
-        .collect();
+    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();
@@ -139,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(
@@ -255,7 +255,8 @@ fn try_find_self_references(
     syntax: &SyntaxNode,
     position: FilePosition,
 ) -> Option<RangeInfo<ReferenceSearchResult>> {
-    let self_token = syntax.token_at_offset(position.offset).find(|t| t.kind() == T![self])?;
+    let FilePosition { file_id, offset } = position;
+    let self_token = syntax.token_at_offset(offset).find(|t| t.kind() == T![self])?;
     let parent = self_token.parent();
     match_ast! {
         match parent {
@@ -276,7 +277,7 @@ fn try_find_self_references(
 
     let declaration = Declaration {
         nav: NavigationTarget {
-            file_id: position.file_id,
+            file_id,
             full_range: self_param.syntax().text_range(),
             focus_range: Some(param_self_token.text_range()),
             name: param_self_token.text().clone(),
@@ -292,7 +293,7 @@ fn try_find_self_references(
             ReferenceAccess::Read
         }),
     };
-    let references = function
+    let refs = function
         .body()
         .map(|body| {
             body.syntax()
@@ -306,14 +307,16 @@ fn try_find_self_references(
                         None
                     }
                 })
-                .map(|token| Reference {
-                    file_range: FileRange { file_id: position.file_id, range: token.text_range() },
+                .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()
         })
         .unwrap_or_default();
+    let mut references = UsageSearchResult::default();
+    references.references.insert(file_id, refs);
 
     Some(RangeInfo::new(
         param_self_token.text_range(),
@@ -1018,12 +1021,14 @@ impl Foo {
             actual += "\n\n";
         }
 
-        for r in &refs.references {
-            format_to!(actual, "{:?} {:?} {:?}", r.file_range.file_id, r.file_range.range, r.kind);
-            if let Some(access) = r.access {
-                format_to!(actual, " {:?}", access);
+        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 {
+                    format_to!(actual, " {:?}", access);
+                }
+                actual += "\n";
             }
-            actual += "\n";
         }
         expect.assert_eq(&actual)
     }
diff --git a/crates/ide/src/references/rename.rs b/crates/ide/src/references/rename.rs
index 3edc43e08..5207388b5 100644
--- a/crates/ide/src/references/rename.rs
+++ b/crates/ide/src/references/rename.rs
@@ -6,9 +6,10 @@ use std::{
 };
 
 use hir::{Module, ModuleDef, ModuleSource, Semantics};
-use ide_db::base_db::{AnchoredPathBuf, FileId, FileRange, SourceDatabaseExt};
 use ide_db::{
+    base_db::{AnchoredPathBuf, FileId, FileRange, SourceDatabaseExt},
     defs::{Definition, NameClass, NameRefClass},
+    search::FileReference,
     RootDatabase,
 };
 use syntax::{
@@ -20,8 +21,8 @@ use test_utils::mark;
 use text_edit::TextEdit;
 
 use crate::{
-    FilePosition, FileSystemEdit, RangeInfo, Reference, ReferenceKind, ReferenceSearchResult,
-    SourceChange, SourceFileEdit, TextRange, TextSize,
+    FilePosition, FileSystemEdit, RangeInfo, ReferenceKind, ReferenceSearchResult, SourceChange,
+    SourceFileEdit, TextRange, TextSize,
 };
 
 type RenameResult<T> = Result<T, RenameError>;
@@ -173,39 +174,46 @@ fn find_all_refs(
         .ok_or_else(|| format_err!("No references found at position"))
 }
 
-fn source_edit_from_reference(
+fn source_edit_from_references(
     sema: &Semantics<RootDatabase>,
-    reference: Reference,
+    file_id: FileId,
+    references: &[FileReference],
     new_name: &str,
 ) -> SourceFileEdit {
-    let mut replacement_text = String::new();
-    let range = match reference.kind {
-        ReferenceKind::FieldShorthandForField => {
-            mark::hit!(test_rename_struct_field_for_shorthand);
-            replacement_text.push_str(new_name);
-            replacement_text.push_str(": ");
-            TextRange::new(reference.file_range.range.start(), reference.file_range.range.start())
-        }
-        ReferenceKind::FieldShorthandForLocal => {
-            mark::hit!(test_rename_local_for_field_shorthand);
-            replacement_text.push_str(": ");
-            replacement_text.push_str(new_name);
-            TextRange::new(reference.file_range.range.end(), reference.file_range.range.end())
-        }
-        ReferenceKind::RecordFieldExprOrPat => {
-            mark::hit!(test_rename_field_expr_pat);
-            replacement_text.push_str(new_name);
-            edit_text_range_for_record_field_expr_or_pat(sema, reference.file_range, new_name)
-        }
-        _ => {
-            replacement_text.push_str(new_name);
-            reference.file_range.range
-        }
-    };
-    SourceFileEdit {
-        file_id: reference.file_range.file_id,
-        edit: TextEdit::replace(range, replacement_text),
+    let mut edit = TextEdit::builder();
+    for reference in references {
+        let mut replacement_text = String::new();
+        let range = match reference.kind {
+            ReferenceKind::FieldShorthandForField => {
+                mark::hit!(test_rename_struct_field_for_shorthand);
+                replacement_text.push_str(new_name);
+                replacement_text.push_str(": ");
+                TextRange::new(reference.range.start(), reference.range.start())
+            }
+            ReferenceKind::FieldShorthandForLocal => {
+                mark::hit!(test_rename_local_for_field_shorthand);
+                replacement_text.push_str(": ");
+                replacement_text.push_str(new_name);
+                TextRange::new(reference.range.end(), reference.range.end())
+            }
+            ReferenceKind::RecordFieldExprOrPat => {
+                mark::hit!(test_rename_field_expr_pat);
+                replacement_text.push_str(new_name);
+                edit_text_range_for_record_field_expr_or_pat(
+                    sema,
+                    FileRange { file_id, range: reference.range },
+                    new_name,
+                )
+            }
+            _ => {
+                replacement_text.push_str(new_name);
+                reference.range
+            }
+        };
+        edit.replace(range, replacement_text);
     }
+
+    SourceFileEdit { file_id, edit: edit.finish() }
 }
 
 fn edit_text_range_for_record_field_expr_or_pat(
@@ -276,10 +284,9 @@ fn rename_mod(
     }
 
     let RangeInfo { range, info: refs } = find_all_refs(sema, position)?;
-    let ref_edits = refs
-        .references
-        .into_iter()
-        .map(|reference| source_edit_from_reference(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)))
@@ -331,17 +338,12 @@ fn rename_to_self(
 
     let RangeInfo { range, info: refs } = find_all_refs(sema, position)?;
 
-    let (param_ref, usages): (Vec<Reference>, Vec<Reference>) = refs
-        .into_iter()
-        .partition(|reference| param_range.intersect(reference.file_range.range).is_some());
-
-    if param_ref.is_empty() {
-        bail!("Parameter to rename not found");
-    }
-
-    let mut edits = usages
-        .into_iter()
-        .map(|reference| source_edit_from_reference(sema, reference, "self"))
+    let mut edits = refs
+        .references()
+        .iter()
+        .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_reference(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)))