Auto merge of #13318 - Veykril:annotations, r=Veykril

Fix annotations not resolving when lens location is set to whole item

Fixes https://github.com/rust-lang/rust-analyzer/issues/13310

5 files changed, 199 insertions, 130 deletions

diff --git a/crates/ide/src/annotations.rs b/crates/ide/src/annotations.rs
index bfbe0db6e4b..f994c284c71 100644
--- a/crates/ide/src/annotations.rs
+++ b/crates/ide/src/annotations.rs
@@ -8,13 +8,15 @@ use ide_db::{
 use syntax::{ast::HasName, AstNode, TextRange};
 
 use crate::{
-    fn_references::find_all_methods,
+    annotations::fn_references::find_all_methods,
     goto_implementation::goto_implementation,
     references::find_all_refs,
     runnables::{runnables, Runnable},
     NavigationTarget, RunnableKind,
 };
 
+mod fn_references;
+
 // Feature: Annotations
 //
 // Provides user with annotations above items for looking up references or impl blocks
@@ -30,8 +32,8 @@ pub struct Annotation {
 #[derive(Debug)]
 pub enum AnnotationKind {
     Runnable(Runnable),
-    HasImpls { file_id: FileId, data: Option<Vec<NavigationTarget>> },
-    HasReferences { file_id: FileId, data: Option<Vec<FileRange>> },
+    HasImpls { pos: FilePosition, data: Option<Vec<NavigationTarget>> },
+    HasReferences { pos: FilePosition, data: Option<Vec<FileRange>> },
 }
 
 pub struct AnnotationConfig {
@@ -68,13 +70,23 @@ pub(crate) fn annotations(
         }
     }
 
+    let mk_ranges = |(range, focus): (_, Option<_>)| {
+        let cmd_target: TextRange = focus.unwrap_or(range);
+        let annotation_range = match config.location {
+            AnnotationLocation::AboveName => cmd_target,
+            AnnotationLocation::AboveWholeItem => range,
+        };
+        let target_pos = FilePosition { file_id, offset: cmd_target.start() };
+        (annotation_range, target_pos)
+    };
+
     visit_file_defs(&Semantics::new(db), file_id, &mut |def| {
         let range = match def {
             Definition::Const(konst) if config.annotate_references => {
-                konst.source(db).and_then(|node| name_range(db, config, node, file_id))
+                konst.source(db).and_then(|node| name_range(db, node, file_id))
             }
             Definition::Trait(trait_) if config.annotate_references || config.annotate_impls => {
-                trait_.source(db).and_then(|node| name_range(db, config, node, file_id))
+                trait_.source(db).and_then(|node| name_range(db, node, file_id))
             }
             Definition::Adt(adt) => match adt {
                 hir::Adt::Enum(enum_) => {
@@ -83,27 +95,29 @@ pub(crate) fn annotations(
                             .variants(db)
                             .into_iter()
                             .map(|variant| {
-                                variant
-                                    .source(db)
-                                    .and_then(|node| name_range(db, config, node, file_id))
+                                variant.source(db).and_then(|node| name_range(db, node, file_id))
                             })
                             .flatten()
                             .for_each(|range| {
+                                let (annotation_range, target_position) = mk_ranges(range);
                                 annotations.push(Annotation {
-                                    range,
-                                    kind: AnnotationKind::HasReferences { file_id, data: None },
+                                    range: annotation_range,
+                                    kind: AnnotationKind::HasReferences {
+                                        pos: target_position,
+                                        data: None,
+                                    },
                                 })
                             })
                     }
                     if config.annotate_references || config.annotate_impls {
-                        enum_.source(db).and_then(|node| name_range(db, config, node, file_id))
+                        enum_.source(db).and_then(|node| name_range(db, node, file_id))
                     } else {
                         None
                     }
                 }
                 _ => {
                     if config.annotate_references || config.annotate_impls {
-                        adt.source(db).and_then(|node| name_range(db, config, node, file_id))
+                        adt.source(db).and_then(|node| name_range(db, node, file_id))
                     } else {
                         None
                     }
@@ -116,33 +130,32 @@ pub(crate) fn annotations(
             Some(range) => range,
             None => return,
         };
-
+        let (annotation_range, target_pos) = mk_ranges(range);
         if config.annotate_impls && !matches!(def, Definition::Const(_)) {
-            annotations
-                .push(Annotation { range, kind: AnnotationKind::HasImpls { file_id, data: None } });
+            annotations.push(Annotation {
+                range: annotation_range,
+                kind: AnnotationKind::HasImpls { pos: target_pos, data: None },
+            });
         }
 
         if config.annotate_references {
             annotations.push(Annotation {
-                range,
-                kind: AnnotationKind::HasReferences { file_id, data: None },
+                range: annotation_range,
+                kind: AnnotationKind::HasReferences { pos: target_pos, data: None },
             });
         }
 
         fn name_range<T: HasName>(
             db: &RootDatabase,
-            config: &AnnotationConfig,
             node: InFile<T>,
             source_file_id: FileId,
-        ) -> Option<TextRange> {
+        ) -> Option<(TextRange, Option<TextRange>)> {
             if let Some(InFile { file_id, value }) = node.original_ast_node(db) {
                 if file_id == source_file_id.into() {
-                    return match config.location {
-                        AnnotationLocation::AboveName => {
-                            value.name().map(|name| name.syntax().text_range())
-                        }
-                        AnnotationLocation::AboveWholeItem => Some(value.syntax().text_range()),
-                    };
+                    return Some((
+                        value.syntax().text_range(),
+                        value.name().map(|name| name.syntax().text_range()),
+                    ));
                 }
             }
             None
@@ -150,12 +163,13 @@ pub(crate) fn annotations(
     });
 
     if config.annotate_method_references {
-        annotations.extend(find_all_methods(db, file_id).into_iter().map(
-            |FileRange { file_id, range }| Annotation {
-                range,
-                kind: AnnotationKind::HasReferences { file_id, data: None },
-            },
-        ));
+        annotations.extend(find_all_methods(db, file_id).into_iter().map(|range| {
+            let (annotation_range, target_range) = mk_ranges(range);
+            Annotation {
+                range: annotation_range,
+                kind: AnnotationKind::HasReferences { pos: target_range, data: None },
+            }
+        }));
     }
 
     annotations
@@ -163,18 +177,11 @@ pub(crate) fn annotations(
 
 pub(crate) fn resolve_annotation(db: &RootDatabase, mut annotation: Annotation) -> Annotation {
     match annotation.kind {
-        AnnotationKind::HasImpls { file_id, ref mut data } => {
-            *data =
-                goto_implementation(db, FilePosition { file_id, offset: annotation.range.start() })
-                    .map(|range| range.info);
+        AnnotationKind::HasImpls { pos, ref mut data } => {
+            *data = goto_implementation(db, pos).map(|range| range.info);
         }
-        AnnotationKind::HasReferences { file_id, ref mut data } => {
-            *data = find_all_refs(
-                &Semantics::new(db),
-                FilePosition { file_id, offset: annotation.range.start() },
-                None,
-            )
-            .map(|result| {
+        AnnotationKind::HasReferences { pos, ref mut data } => {
+            *data = find_all_refs(&Semantics::new(db), pos, None).map(|result| {
                 result
                     .into_iter()
                     .flat_map(|res| res.references)
@@ -268,9 +275,12 @@ fn main() {
                     Annotation {
                         range: 6..10,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 6,
+                            },
                             data: Some(
                                 [
                                     FileRange {
@@ -286,9 +296,12 @@ fn main() {
                     Annotation {
                         range: 30..36,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 30,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -297,9 +310,12 @@ fn main() {
                     Annotation {
                         range: 53..57,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 53,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -344,9 +360,12 @@ fn main() {
                     Annotation {
                         range: 7..11,
                         kind: HasImpls {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 7,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -355,9 +374,12 @@ fn main() {
                     Annotation {
                         range: 7..11,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 7,
+                            },
                             data: Some(
                                 [
                                     FileRange {
@@ -373,9 +395,12 @@ fn main() {
                     Annotation {
                         range: 17..21,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 17,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -424,9 +449,12 @@ fn main() {
                     Annotation {
                         range: 7..11,
                         kind: HasImpls {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 7,
+                            },
                             data: Some(
                                 [
                                     NavigationTarget {
@@ -445,9 +473,12 @@ fn main() {
                     Annotation {
                         range: 7..11,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 7,
+                            },
                             data: Some(
                                 [
                                     FileRange {
@@ -469,9 +500,12 @@ fn main() {
                     Annotation {
                         range: 20..31,
                         kind: HasImpls {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 20,
+                            },
                             data: Some(
                                 [
                                     NavigationTarget {
@@ -490,9 +524,12 @@ fn main() {
                     Annotation {
                         range: 20..31,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 20,
+                            },
                             data: Some(
                                 [
                                     FileRange {
@@ -508,9 +545,12 @@ fn main() {
                     Annotation {
                         range: 69..73,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 69,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -551,9 +591,12 @@ fn main() {}
                     Annotation {
                         range: 3..7,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 3,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -602,9 +645,12 @@ fn main() {
                     Annotation {
                         range: 7..11,
                         kind: HasImpls {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 7,
+                            },
                             data: Some(
                                 [
                                     NavigationTarget {
@@ -623,9 +669,12 @@ fn main() {
                     Annotation {
                         range: 7..11,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 7,
+                            },
                             data: Some(
                                 [
                                     FileRange {
@@ -647,9 +696,12 @@ fn main() {
                     Annotation {
                         range: 33..44,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 33,
+                            },
                             data: Some(
                                 [
                                     FileRange {
@@ -665,9 +717,12 @@ fn main() {
                     Annotation {
                         range: 61..65,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 61,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -761,9 +816,12 @@ mod tests {
                     Annotation {
                         range: 3..7,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 3,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -821,9 +879,12 @@ struct Foo;
                     Annotation {
                         range: 0..71,
                         kind: HasImpls {
-                            file_id: FileId(
-                                0,
-                            ),
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 67,
+                            },
                             data: Some(
                                 [],
                             ),
@@ -832,10 +893,15 @@ struct Foo;
                     Annotation {
                         range: 0..71,
                         kind: HasReferences {
-                            file_id: FileId(
-                                0,
+                            pos: FilePosition {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 67,
+                            },
+                            data: Some(
+                                [],
                             ),
-                            data: None,
                         },
                     },
                 ]
diff --git a/crates/ide/src/fn_references.rs b/crates/ide/src/annotations/fn_references.rs
index 63fb322cea0..0cadf125fec 100644
--- a/crates/ide/src/fn_references.rs
+++ b/crates/ide/src/annotations/fn_references.rs
@@ -4,30 +4,38 @@
 use hir::Semantics;
 use ide_assists::utils::test_related_attribute;
 use ide_db::RootDatabase;
-use syntax::{ast, ast::HasName, AstNode, SyntaxNode};
+use syntax::{ast, ast::HasName, AstNode, SyntaxNode, TextRange};
 
-use crate::{FileId, FileRange};
+use crate::FileId;
 
-pub(crate) fn find_all_methods(db: &RootDatabase, file_id: FileId) -> Vec<FileRange> {
+pub(super) fn find_all_methods(
+    db: &RootDatabase,
+    file_id: FileId,
+) -> Vec<(TextRange, Option<TextRange>)> {
     let sema = Semantics::new(db);
     let source_file = sema.parse(file_id);
-    source_file.syntax().descendants().filter_map(|it| method_range(it, file_id)).collect()
+    source_file.syntax().descendants().filter_map(|it| method_range(it)).collect()
 }
 
-fn method_range(item: SyntaxNode, file_id: FileId) -> Option<FileRange> {
+fn method_range(item: SyntaxNode) -> Option<(TextRange, Option<TextRange>)> {
     ast::Fn::cast(item).and_then(|fn_def| {
         if test_related_attribute(&fn_def).is_some() {
             None
         } else {
-            fn_def.name().map(|name| FileRange { file_id, range: name.syntax().text_range() })
+            Some((
+                fn_def.syntax().text_range(),
+                fn_def.name().map(|name| name.syntax().text_range()),
+            ))
         }
     })
 }
 
 #[cfg(test)]
 mod tests {
+    use syntax::TextRange;
+
     use crate::fixture;
-    use crate::{FileRange, TextSize};
+    use crate::TextSize;
     use std::ops::RangeInclusive;
 
     #[test]
@@ -42,7 +50,7 @@ mod tests {
         "#,
         );
 
-        let refs = analysis.find_all_methods(pos.file_id).unwrap();
+        let refs = super::find_all_methods(&analysis.db, pos.file_id);
         check_result(&refs, &[3..=13, 27..=33, 47..=57]);
     }
 
@@ -57,7 +65,7 @@ mod tests {
         "#,
         );
 
-        let refs = analysis.find_all_methods(pos.file_id).unwrap();
+        let refs = super::find_all_methods(&analysis.db, pos.file_id);
         check_result(&refs, &[19..=22, 35..=38]);
     }
 
@@ -78,17 +86,18 @@ mod tests {
         "#,
         );
 
-        let refs = analysis.find_all_methods(pos.file_id).unwrap();
+        let refs = super::find_all_methods(&analysis.db, pos.file_id);
         check_result(&refs, &[28..=34]);
     }
 
-    fn check_result(refs: &[FileRange], expected: &[RangeInclusive<u32>]) {
+    fn check_result(refs: &[(TextRange, Option<TextRange>)], expected: &[RangeInclusive<u32>]) {
         assert_eq!(refs.len(), expected.len());
 
-        for (i, item) in refs.iter().enumerate() {
+        for (i, &(full, focus)) in refs.iter().enumerate() {
             let range = &expected[i];
-            assert_eq!(TextSize::from(*range.start()), item.range.start());
-            assert_eq!(TextSize::from(*range.end()), item.range.end());
+            let item = focus.unwrap_or(full);
+            assert_eq!(TextSize::from(*range.start()), item.start());
+            assert_eq!(TextSize::from(*range.end()), item.end());
         }
     }
 }
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 7820b67142f..77fe0dbf556 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -31,7 +31,6 @@ mod highlight_related;
 mod expand_macro;
 mod extend_selection;
 mod file_structure;
-mod fn_references;
 mod folding_ranges;
 mod goto_declaration;
 mod goto_definition;
@@ -429,11 +428,6 @@ impl Analysis {
         self.with_db(|db| references::find_all_refs(&Semantics::new(db), position, search_scope))
     }
 
-    /// Finds all methods and free functions for the file. Does not return tests!
-    pub fn find_all_methods(&self, file_id: FileId) -> Cancellable<Vec<FileRange>> {
-        self.with_db(|db| fn_references::find_all_methods(db, file_id))
-    }
-
     /// Returns a short text describing element at position.
     pub fn hover(
         &self,
diff --git a/crates/rust-analyzer/src/from_proto.rs b/crates/rust-analyzer/src/from_proto.rs
index 7bdd34d1f09..f2db9a27334 100644
--- a/crates/rust-analyzer/src/from_proto.rs
+++ b/crates/rust-analyzer/src/from_proto.rs
@@ -95,22 +95,22 @@ pub(crate) fn annotation(
 
     match resolve {
         lsp_ext::CodeLensResolveData::Impls(params) => {
-            let file_id =
-                snap.url_to_file_id(&params.text_document_position_params.text_document.uri)?;
+            let pos @ FilePosition { file_id, .. } =
+                file_position(snap, params.text_document_position_params)?;
             let line_index = snap.file_line_index(file_id)?;
 
             Ok(Annotation {
                 range: text_range(&line_index, code_lens.range)?,
-                kind: AnnotationKind::HasImpls { file_id, data: None },
+                kind: AnnotationKind::HasImpls { pos, data: None },
             })
         }
         lsp_ext::CodeLensResolveData::References(params) => {
-            let file_id = snap.url_to_file_id(&params.text_document.uri)?;
+            let pos @ FilePosition { file_id, .. } = file_position(snap, params)?;
             let line_index = snap.file_line_index(file_id)?;
 
             Ok(Annotation {
                 range: text_range(&line_index, code_lens.range)?,
-                kind: AnnotationKind::HasReferences { file_id, data: None },
+                kind: AnnotationKind::HasReferences { pos, data: None },
             })
         }
     }
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index 1de801e23e5..5936454a7c5 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -1177,13 +1177,13 @@ pub(crate) fn code_lens(
                 })
             }
         }
-        AnnotationKind::HasImpls { file_id, data } => {
+        AnnotationKind::HasImpls { pos: file_range, data } => {
             if !client_commands_config.show_reference {
                 return Ok(());
             }
-            let line_index = snap.file_line_index(file_id)?;
+            let line_index = snap.file_line_index(file_range.file_id)?;
             let annotation_range = range(&line_index, annotation.range);
-            let url = url(snap, file_id);
+            let url = url(snap, file_range.file_id);
 
             let id = lsp_types::TextDocumentIdentifier { uri: url.clone() };
 
@@ -1221,13 +1221,13 @@ pub(crate) fn code_lens(
                 data: Some(to_value(lsp_ext::CodeLensResolveData::Impls(goto_params)).unwrap()),
             })
         }
-        AnnotationKind::HasReferences { file_id, data } => {
+        AnnotationKind::HasReferences { pos: file_range, data } => {
             if !client_commands_config.show_reference {
                 return Ok(());
             }
-            let line_index = snap.file_line_index(file_id)?;
+            let line_index = snap.file_line_index(file_range.file_id)?;
             let annotation_range = range(&line_index, annotation.range);
-            let url = url(snap, file_id);
+            let url = url(snap, file_range.file_id);
 
             let id = lsp_types::TextDocumentIdentifier { uri: url.clone() };