1 files changed, 112 insertions, 52 deletions

diff --git a/crates/ide_db/src/search.rs b/crates/ide_db/src/search.rs
index ff10f71c3..b5fa46642 100644
--- a/crates/ide_db/src/search.rs
+++ b/crates/ide_db/src/search.rs
@@ -18,9 +18,43 @@ use crate::{
     RootDatabase,
 };
 
+#[derive(Debug, Default, Clone)]
+pub struct UsageSearchResult {
+    pub references: FxHashMap<FileId, Vec<FileReference>>,
+}
+
+impl UsageSearchResult {
+    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().flat_map(|(&file_id, refs)| {
+            refs.iter().map(move |&FileReference { range, .. }| FileRange { file_id, range })
+        })
+    }
+}
+
+impl IntoIterator for UsageSearchResult {
+    type Item = (FileId, Vec<FileReference>);
+    type IntoIter = <FxHashMap<FileId, Vec<FileReference>> as IntoIterator>::IntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.references.into_iter()
+    }
+}
+
 #[derive(Debug, Clone)]
-pub struct Reference {
-    pub file_range: FileRange,
+pub struct FileReference {
+    pub range: TextRange,
     pub kind: ReferenceKind,
     pub access: Option<ReferenceAccess>,
 }
@@ -121,31 +155,55 @@ impl Definition {
 
         if let Definition::Local(var) = self {
             let range = match var.parent(db) {
-                DefWithBody::Function(f) => f.source(db).value.syntax().text_range(),
-                DefWithBody::Const(c) => c.source(db).value.syntax().text_range(),
-                DefWithBody::Static(s) => s.source(db).value.syntax().text_range(),
+                DefWithBody::Function(f) => {
+                    f.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
+                DefWithBody::Const(c) => {
+                    c.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
+                DefWithBody::Static(s) => {
+                    s.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
             };
             let mut res = FxHashMap::default();
-            res.insert(file_id, Some(range));
+            res.insert(file_id, range);
             return SearchScope::new(res);
         }
 
-        if let Definition::LifetimeParam(param) = self {
+        if let Definition::GenericParam(hir::GenericParam::LifetimeParam(param)) = self {
             let range = match param.parent(db) {
-                hir::GenericDef::Function(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::Function(it) => {
+                    it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
                 hir::GenericDef::Adt(it) => match it {
-                    hir::Adt::Struct(it) => it.source(db).value.syntax().text_range(),
-                    hir::Adt::Union(it) => it.source(db).value.syntax().text_range(),
-                    hir::Adt::Enum(it) => it.source(db).value.syntax().text_range(),
+                    hir::Adt::Struct(it) => {
+                        it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                    }
+                    hir::Adt::Union(it) => {
+                        it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                    }
+                    hir::Adt::Enum(it) => {
+                        it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                    }
                 },
-                hir::GenericDef::Trait(it) => it.source(db).value.syntax().text_range(),
-                hir::GenericDef::TypeAlias(it) => it.source(db).value.syntax().text_range(),
-                hir::GenericDef::Impl(it) => it.source(db).value.syntax().text_range(),
-                hir::GenericDef::Variant(it) => it.source(db).value.syntax().text_range(),
-                hir::GenericDef::Const(it) => it.source(db).value.syntax().text_range(),
+                hir::GenericDef::Trait(it) => {
+                    it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
+                hir::GenericDef::TypeAlias(it) => {
+                    it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
+                hir::GenericDef::Impl(it) => {
+                    it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
+                hir::GenericDef::Variant(it) => {
+                    it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
+                hir::GenericDef::Const(it) => {
+                    it.source(db).and_then(|src| Some(src.value.syntax().text_range()))
+                }
             };
             let mut res = FxHashMap::default();
-            res.insert(file_id, Some(range));
+            res.insert(file_id, range);
             return SearchScope::new(res);
         }
 
@@ -228,23 +286,23 @@ impl<'a> FindUsages<'a> {
 
     pub fn at_least_one(self) -> bool {
         let mut found = false;
-        self.search(&mut |_reference| {
+        self.search(&mut |_, _| {
             found = true;
             true
         });
         found
     }
 
-    pub fn all(self) -> Vec<Reference> {
-        let mut res = Vec::new();
-        self.search(&mut |reference| {
-            res.push(reference);
+    pub fn all(self) -> UsageSearchResult {
+        let mut res = UsageSearchResult::default();
+        self.search(&mut |file_id, reference| {
+            res.references.entry(file_id).or_default().push(reference);
             false
         });
         res
     }
 
-    fn search(self, sink: &mut dyn FnMut(Reference) -> bool) {
+    fn search(self, sink: &mut dyn FnMut(FileId, FileReference) -> bool) {
         let _p = profile::span("FindUsages:search");
         let sema = self.sema;
 
@@ -296,16 +354,14 @@ impl<'a> FindUsages<'a> {
     fn found_lifetime(
         &self,
         lifetime: &ast::Lifetime,
-        sink: &mut dyn FnMut(Reference) -> bool,
+        sink: &mut dyn FnMut(FileId, FileReference) -> bool,
     ) -> bool {
         match NameRefClass::classify_lifetime(self.sema, lifetime) {
             Some(NameRefClass::Definition(def)) if &def == self.def => {
-                let reference = Reference {
-                    file_range: self.sema.original_range(lifetime.syntax()),
-                    kind: ReferenceKind::Lifetime,
-                    access: None,
-                };
-                sink(reference)
+                let FileRange { file_id, range } = self.sema.original_range(lifetime.syntax());
+                let reference =
+                    FileReference { range, kind: ReferenceKind::Lifetime, access: None };
+                sink(file_id, reference)
             }
             _ => false, // not a usage
         }
@@ -314,7 +370,7 @@ impl<'a> FindUsages<'a> {
     fn found_name_ref(
         &self,
         name_ref: &ast::NameRef,
-        sink: &mut dyn FnMut(Reference) -> bool,
+        sink: &mut dyn FnMut(FileId, FileReference) -> bool,
     ) -> bool {
         match NameRefClass::classify(self.sema, &name_ref) {
             Some(NameRefClass::Definition(def)) if &def == self.def => {
@@ -328,46 +384,50 @@ impl<'a> FindUsages<'a> {
                     ReferenceKind::Other
                 };
 
-                let reference = Reference {
-                    file_range: self.sema.original_range(name_ref.syntax()),
-                    kind,
-                    access: reference_access(&def, &name_ref),
-                };
-                sink(reference)
+                let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
+                let reference =
+                    FileReference { range, kind, access: reference_access(&def, &name_ref) };
+                sink(file_id, reference)
             }
             Some(NameRefClass::FieldShorthand { local_ref: local, field_ref: field }) => {
+                let FileRange { file_id, range } = self.sema.original_range(name_ref.syntax());
                 let reference = match self.def {
-                    Definition::Field(_) if &field == self.def => Reference {
-                        file_range: self.sema.original_range(name_ref.syntax()),
+                    Definition::Field(_) if &field == self.def => FileReference {
+                        range,
                         kind: ReferenceKind::FieldShorthandForField,
                         access: reference_access(&field, &name_ref),
                     },
-                    Definition::Local(l) if &local == l => Reference {
-                        file_range: self.sema.original_range(name_ref.syntax()),
+                    Definition::Local(l) if &local == l => FileReference {
+                        range,
                         kind: ReferenceKind::FieldShorthandForLocal,
                         access: reference_access(&Definition::Local(local), &name_ref),
                     },
                     _ => return false, // not a usage
                 };
-                sink(reference)
+                sink(file_id, reference)
             }
             _ => false, // not a usage
         }
     }
 
-    fn found_name(&self, name: &ast::Name, sink: &mut dyn FnMut(Reference) -> bool) -> bool {
+    fn found_name(
+        &self,
+        name: &ast::Name,
+        sink: &mut dyn FnMut(FileId, FileReference) -> bool,
+    ) -> bool {
         match NameClass::classify(self.sema, name) {
             Some(NameClass::PatFieldShorthand { local_def: _, field_ref }) => {
-                let reference = match self.def {
-                    Definition::Field(_) if &field_ref == self.def => Reference {
-                        file_range: self.sema.original_range(name.syntax()),
-                        kind: ReferenceKind::FieldShorthandForField,
-                        // FIXME: mutable patterns should have `Write` access
-                        access: Some(ReferenceAccess::Read),
-                    },
-                    _ => return false, // not a usage
+                if !matches!(self.def, Definition::Field(_) if &field_ref == self.def) {
+                    return false;
+                }
+                let FileRange { file_id, range } = self.sema.original_range(name.syntax());
+                let reference = FileReference {
+                    range,
+                    kind: ReferenceKind::FieldShorthandForField,
+                    // FIXME: mutable patterns should have `Write` access
+                    access: Some(ReferenceAccess::Read),
                 };
-                sink(reference)
+                sink(file_id, reference)
             }
             _ => false, // not a usage
         }