Merge #10103

10103: fix: make "find references" multi-token mapping aware r=jonas-schievink a=jonas-schievink

Part of https://github.com/rust-analyzer/rust-analyzer/issues/10070

I was hoping that this would fix "find references" on salsa queries, but salsa emits multiple defs sharing the same span (the trait method, and an impl of that trait).

Co-authored-by: Jonas Schievink <jonasschievink@gmail.com>

crates/ide/src/annotations.rs

@@ -134,8 +134,8 @@ pub(crate) fn resolve_annotation(db: &RootDatabase, mut annotation: Annotation)
         AnnotationKind::HasReferences { position, data } => {
             *data = find_all_refs(&Semantics::new(db), *position, None).map(|result| {
                 result
-                    .references
                     .into_iter()
+                    .flat_map(|res| res.references)
                     .map(|(file_id, access)| {
                         access.into_iter().map(move |(range, _)| FileRange { file_id, range })
                     })

crates/ide/src/call_hierarchy.rs

@@ -46,7 +46,7 @@ pub(crate) fn incoming_calls(db: &RootDatabase, position: FilePosition) -> Optio
 
     let mut calls = CallLocations::default();
 
-    for (file_id, references) in refs.references {
+    for (file_id, references) in refs.into_iter().flat_map(|refs| refs.references) {
         let file = sema.parse(file_id);
         let file = file.syntax();
         for (relative_range, token) in references

crates/ide/src/lib.rs

@@ -405,7 +405,7 @@ impl Analysis {
         &self,
         position: FilePosition,
         search_scope: Option<SearchScope>,
-    ) -> Cancellable<Option<ReferenceSearchResult>> {
+    ) -> Cancellable<Option<Vec<ReferenceSearchResult>>> {
         self.with_db(|db| references::find_all_refs(&Semantics::new(db), position, search_scope))
     }
 

crates/ide/src/references.rs

@@ -9,6 +9,9 @@
 //! at the index that the match starts at and its tree parent is
 //! resolved to the search element definition, we get a reference.
 
+use std::iter;
+
+use either::Either;
 use hir::{PathResolution, Semantics};
 use ide_db::{
     base_db::FileId,
@@ -52,76 +55,91 @@ pub(crate) fn find_all_refs(
     sema: &Semantics<RootDatabase>,
     position: FilePosition,
     search_scope: Option<SearchScope>,
-) -> Option<ReferenceSearchResult> {
+) -> Option<Vec<ReferenceSearchResult>> {
     let _p = profile::span("find_all_refs");
     let syntax = sema.parse(position.file_id).syntax().clone();
 
     let mut is_literal_search = false;
-    let def = if let Some(name) = name_for_constructor_search(&syntax, position) {
-        is_literal_search = true;
-        match NameClass::classify(sema, &name)? {
-            NameClass::Definition(it) | NameClass::ConstReference(it) => it,
-            NameClass::PatFieldShorthand { local_def: _, field_ref } => {
-                Definition::Field(field_ref)
-            }
+    let defs = match name_for_constructor_search(&syntax, position) {
+        Some(name) => {
+            is_literal_search = true;
+            let def = match NameClass::classify(sema, &name)? {
+                NameClass::Definition(it) | NameClass::ConstReference(it) => it,
+                NameClass::PatFieldShorthand { local_def: _, field_ref } => {
+                    Definition::Field(field_ref)
+                }
+            };
+            Either::Left(iter::once(def))
         }
-    } else {
-        find_def(sema, &syntax, position.offset)?
+        None => Either::Right(find_defs(sema, &syntax, position.offset)),
     };
 
-    let mut usages = def.usages(sema).set_scope(search_scope).include_self_refs().all();
-    let declaration = match def {
-        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
-            Some(NavigationTarget::from_module_to_decl(sema.db, module))
-        }
-        def => def.try_to_nav(sema.db),
-    }
-    .map(|nav| {
-        let decl_range = nav.focus_or_full_range();
-        Declaration { nav, access: decl_access(&def, &syntax, decl_range) }
-    });
-    if is_literal_search {
-        retain_adt_literal_usages(&mut usages, def, sema);
-    }
+    Some(
+        defs.into_iter()
+            .map(|def| {
+                let mut usages =
+                    def.usages(sema).set_scope(search_scope.clone()).include_self_refs().all();
+                let declaration = match def {
+                    Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
+                        Some(NavigationTarget::from_module_to_decl(sema.db, module))
+                    }
+                    def => def.try_to_nav(sema.db),
+                }
+                .map(|nav| {
+                    let decl_range = nav.focus_or_full_range();
+                    Declaration { nav, access: decl_access(&def, &syntax, decl_range) }
+                });
+                if is_literal_search {
+                    retain_adt_literal_usages(&mut usages, def, sema);
+                }
 
-    let references = usages
-        .into_iter()
-        .map(|(file_id, refs)| {
-            (file_id, refs.into_iter().map(|file_ref| (file_ref.range, file_ref.access)).collect())
-        })
-        .collect();
+                let references = usages
+                    .into_iter()
+                    .map(|(file_id, refs)| {
+                        (
+                            file_id,
+                            refs.into_iter()
+                                .map(|file_ref| (file_ref.range, file_ref.access))
+                                .collect(),
+                        )
+                    })
+                    .collect();
 
-    Some(ReferenceSearchResult { declaration, references })
+                ReferenceSearchResult { declaration, references }
+            })
+            .collect(),
+    )
 }
 
-pub(crate) fn find_def(
-    sema: &Semantics<RootDatabase>,
+pub(crate) fn find_defs<'a>(
+    sema: &'a Semantics<RootDatabase>,
     syntax: &SyntaxNode,
     offset: TextSize,
-) -> Option<Definition> {
-    let def = match sema.find_node_at_offset_with_descend(syntax, offset)? {
-        ast::NameLike::NameRef(name_ref) => match NameRefClass::classify(sema, &name_ref)? {
-            NameRefClass::Definition(def) => def,
-            NameRefClass::FieldShorthand { local_ref, field_ref: _ } => {
-                Definition::Local(local_ref)
-            }
-        },
-        ast::NameLike::Name(name) => match NameClass::classify(sema, &name)? {
-            NameClass::Definition(it) | NameClass::ConstReference(it) => it,
-            NameClass::PatFieldShorthand { local_def, field_ref: _ } => {
-                Definition::Local(local_def)
-            }
-        },
-        ast::NameLike::Lifetime(lifetime) => NameRefClass::classify_lifetime(sema, &lifetime)
-            .and_then(|class| match class {
-                NameRefClass::Definition(it) => Some(it),
-                _ => None,
-            })
-            .or_else(|| {
-                NameClass::classify_lifetime(sema, &lifetime).and_then(NameClass::defined)
-            })?,
-    };
-    Some(def)
+) -> impl Iterator<Item = Definition> + 'a {
+    sema.find_nodes_at_offset_with_descend(syntax, offset).filter_map(move |node| {
+        Some(match node {
+            ast::NameLike::NameRef(name_ref) => match NameRefClass::classify(sema, &name_ref)? {
+                NameRefClass::Definition(def) => def,
+                NameRefClass::FieldShorthand { local_ref, field_ref: _ } => {
+                    Definition::Local(local_ref)
+                }
+            },
+            ast::NameLike::Name(name) => match NameClass::classify(sema, &name)? {
+                NameClass::Definition(it) | NameClass::ConstReference(it) => it,
+                NameClass::PatFieldShorthand { local_def, field_ref: _ } => {
+                    Definition::Local(local_def)
+                }
+            },
+            ast::NameLike::Lifetime(lifetime) => NameRefClass::classify_lifetime(sema, &lifetime)
+                .and_then(|class| match class {
+                    NameRefClass::Definition(it) => Some(it),
+                    _ => None,
+                })
+                .or_else(|| {
+                    NameClass::classify_lifetime(sema, &lifetime).and_then(NameClass::defined)
+                })?,
+        })
+    })
 }
 
 pub(crate) fn decl_access(
@@ -609,6 +627,7 @@ impl Foo {
             expect![[r#"
                 f Function FileId(0) 27..43 30..31
 
+                (no references)
             "#]],
         );
     }
@@ -626,6 +645,7 @@ enum Foo {
             expect![[r#"
                 B Variant FileId(0) 22..23 22..23
 
+                (no references)
             "#]],
         );
     }
@@ -643,6 +663,7 @@ enum Foo {
             expect![[r#"
                 field Field FileId(0) 26..35 26..31
 
+                (no references)
             "#]],
         );
     }
@@ -744,6 +765,7 @@ use self$0;
             expect![[r#"
                 Module FileId(0) 0..10
 
+                (no references)
             "#]],
         );
     }
@@ -1065,21 +1087,29 @@ impl Foo {
         let refs = analysis.find_all_refs(pos, search_scope).unwrap().unwrap();
 
         let mut actual = String::new();
-        if let Some(decl) = refs.declaration {
-            format_to!(actual, "{}", decl.nav.debug_render());
-            if let Some(access) = decl.access {
-                format_to!(actual, " {:?}", access)
-            }
+        for refs in refs {
             actual += "\n\n";
-        }
 
-        for (file_id, references) in refs.references {
-            for (range, access) in references {
-                format_to!(actual, "{:?} {:?}", file_id, range);
-                if let Some(access) = access {
-                    format_to!(actual, " {:?}", access);
+            if let Some(decl) = refs.declaration {
+                format_to!(actual, "{}", decl.nav.debug_render());
+                if let Some(access) = decl.access {
+                    format_to!(actual, " {:?}", access)
                 }
-                actual += "\n";
+                actual += "\n\n";
+            }
+
+            for (file_id, references) in &refs.references {
+                for (range, access) in references {
+                    format_to!(actual, "{:?} {:?}", file_id, range);
+                    if let Some(access) = access {
+                        format_to!(actual, " {:?}", access);
+                    }
+                    actual += "\n";
+                }
+            }
+
+            if refs.references.is_empty() {
+                actual += "(no references)\n";
             }
         }
         expect.assert_eq(actual.trim_start())
@@ -1440,4 +1470,38 @@ m$0!();
             "#]],
         );
     }
+
+    #[test]
+    fn multi_def() {
+        check(
+            r#"
+macro_rules! m {
+    ($name:ident) => {
+        mod module {
+            pub fn $name() {}
+        }
+
+        pub fn $name() {}
+    }
+}
+
+m!(func$0);
+
+fn f() {
+    func();
+    module::func();
+}
+            "#,
+            expect![[r#"
+                func Function FileId(0) 137..146 140..144
+
+                FileId(0) 161..165
+
+
+                func Function FileId(0) 137..146 140..144
+
+                FileId(0) 181..185
+            "#]],
+        )
+    }
 }

crates/ide/src/runnables.rs

@@ -226,7 +226,7 @@ fn find_related_tests(
     tests: &mut FxHashSet<Runnable>,
 ) {
     if let Some(refs) = references::find_all_refs(sema, position, search_scope) {
-        for (file_id, refs) in refs.references {
+        for (file_id, refs) in refs.into_iter().flat_map(|refs| refs.references) {
             let file = sema.parse(file_id);
             let file = file.syntax();
             let functions = refs.iter().filter_map(|(range, _)| {

crates/ide_db/src/search.rs

@@ -71,6 +71,7 @@ pub enum ReferenceAccess {
 /// For `pub(crate)` things it's a crate, for `pub` things it's a crate and dependant crates.
 /// In some cases, the location of the references is known to within a `TextRange`,
 /// e.g. for things like local variables.
+#[derive(Clone)]
 pub struct SearchScope {
     entries: FxHashMap<FileId, Option<TextRange>>,
 }

crates/rust-analyzer/src/handlers.rs

@@ -930,21 +930,25 @@ pub(crate) fn handle_references(
         Some(refs) => refs,
     };
 
-    let decl = if params.context.include_declaration {
-        refs.declaration.map(|decl| FileRange {
-            file_id: decl.nav.file_id,
-            range: decl.nav.focus_or_full_range(),
-        })
-    } else {
-        None
-    };
+    let include_declaration = params.context.include_declaration;
     let locations = refs
-        .references
         .into_iter()
-        .flat_map(|(file_id, refs)| {
-            refs.into_iter().map(move |(range, _)| FileRange { file_id, range })
+        .flat_map(|refs| {
+            let decl = if include_declaration {
+                refs.declaration.map(|decl| FileRange {
+                    file_id: decl.nav.file_id,
+                    range: decl.nav.focus_or_full_range(),
+                })
+            } else {
+                None
+            };
+            refs.references
+                .into_iter()
+                .flat_map(|(file_id, refs)| {
+                    refs.into_iter().map(move |(range, _)| FileRange { file_id, range })
+                })
+                .chain(decl)
         })
-        .chain(decl)
         .filter_map(|frange| to_proto::location(&snap, frange).ok())
         .collect();
 
@@ -1515,8 +1519,8 @@ fn show_ref_command_link(
             let line_index = snap.file_line_index(position.file_id).ok()?;
             let position = to_proto::position(&line_index, position.offset);
             let locations: Vec<_> = ref_search_res
-                .references
                 .into_iter()
+                .flat_map(|res| res.references)
                 .flat_map(|(file_id, ranges)| {
                     ranges.into_iter().filter_map(move |(range, _)| {
                         to_proto::location(snap, FileRange { file_id, range }).ok()