Add method references CodeLens

crates/ide/src/fn_references.rs

@@ -0,0 +1,95 @@
+//! This module implements a methods and free functions search in the specified file.
+//! We have to skip tests, so cannot reuse file_structure module.
+
+use hir::Semantics;
+use ide_db::RootDatabase;
+use syntax::{ast, ast::NameOwner, AstNode, SyntaxNode};
+
+use crate::{runnables::has_test_related_attribute, FileId, FileRange};
+
+pub(crate) fn find_all_methods(db: &RootDatabase, file_id: FileId) -> Vec<FileRange> {
+    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()
+}
+
+fn method_range(item: SyntaxNode, file_id: FileId) -> Option<FileRange> {
+    ast::Fn::cast(item).and_then(|fn_def| {
+        if has_test_related_attribute(&fn_def) {
+            None
+        } else {
+            fn_def.name().map(|name| FileRange { file_id, range: name.syntax().text_range() })
+        }
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::mock_analysis::analysis_and_position;
+    use crate::{FileRange, TextSize};
+    use std::ops::RangeInclusive;
+
+    #[test]
+    fn test_find_all_methods() {
+        let (analysis, pos) = analysis_and_position(
+            r#"
+            //- /lib.rs
+            fn private_fn() {<|>}
+
+            pub fn pub_fn() {}
+
+            pub fn generic_fn<T>(arg: T) {}
+        "#,
+        );
+
+        let refs = analysis.find_all_methods(pos.file_id).unwrap();
+        check_result(&refs, &[3..=13, 27..=33, 47..=57]);
+    }
+
+    #[test]
+    fn test_find_trait_methods() {
+        let (analysis, pos) = analysis_and_position(
+            r#"
+            //- /lib.rs
+            trait Foo {
+                fn bar() {<|>}
+                fn baz() {}
+            }
+        "#,
+        );
+
+        let refs = analysis.find_all_methods(pos.file_id).unwrap();
+        check_result(&refs, &[19..=22, 35..=38]);
+    }
+
+    #[test]
+    fn test_skip_tests() {
+        let (analysis, pos) = analysis_and_position(
+            r#"
+            //- /lib.rs
+            #[test]
+            fn foo() {<|>}
+
+            pub fn pub_fn() {}
+
+            mod tests {
+                #[test]
+                fn bar() {}
+            }
+        "#,
+        );
+
+        let refs = analysis.find_all_methods(pos.file_id).unwrap();
+        check_result(&refs, &[28..=34]);
+    }
+
+    fn check_result(refs: &[FileRange], expected: &[RangeInclusive<u32>]) {
+        assert_eq!(refs.len(), expected.len());
+
+        for (i, item) 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());
+        }
+    }
+}

crates/ide/src/lib.rs

@@ -38,6 +38,7 @@ mod join_lines;
 mod matching_brace;
 mod parent_module;
 mod references;
+mod fn_references;
 mod runnables;
 mod status;
 mod syntax_highlighting;
@@ -369,6 +370,11 @@ impl Analysis {
         })
     }
 
+    /// Finds all methods and free functions for the file. Does not return tests!
+    pub fn find_all_methods(&self, file_id: FileId) -> Cancelable<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,

crates/ide/src/runnables.rs

@@ -203,7 +203,7 @@ impl TestAttr {
 ///
 /// It may produce false positives, for example, `#[wasm_bindgen_test]` requires a different command to run the test,
 /// but it's better than not to have the runnables for the tests at all.
-fn has_test_related_attribute(fn_def: &ast::Fn) -> bool {
+pub(crate) fn has_test_related_attribute(fn_def: &ast::Fn) -> bool {
     fn_def
         .attrs()
         .filter_map(|attr| attr.path())

crates/rust-analyzer/src/config.rs

@@ -74,19 +74,18 @@ pub struct LensConfig {
     pub run: bool,
     pub debug: bool,
     pub implementations: bool,
+    pub method_refs: bool,
 }
 
 impl Default for LensConfig {
     fn default() -> Self {
-        Self { run: true, debug: true, implementations: true }
+        Self { run: true, debug: true, implementations: true, method_refs: false }
     }
 }
 
 impl LensConfig {
-    pub const NO_LENS: LensConfig = Self { run: false, debug: false, implementations: false };
-
     pub fn any(&self) -> bool {
-        self.implementations || self.runnable()
+        self.implementations || self.runnable() || self.references()
     }
 
     pub fn none(&self) -> bool {
@@ -96,6 +95,10 @@ impl LensConfig {
     pub fn runnable(&self) -> bool {
         self.run || self.debug
     }
+
+    pub fn references(&self) -> bool {
+        self.method_refs
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -278,6 +281,7 @@ impl Config {
             run: data.lens_enable && data.lens_run,
             debug: data.lens_enable && data.lens_debug,
             implementations: data.lens_enable && data.lens_implementations,
+            method_refs: data.lens_enable && data.lens_methodReferences,
         };
 
         if !data.linkedProjects.is_empty() {
@@ -459,10 +463,11 @@ config_data! {
         inlayHints_parameterHints: bool     = true,
         inlayHints_typeHints: bool          = true,
 
-        lens_debug: bool           = true,
-        lens_enable: bool          = true,
-        lens_implementations: bool = true,
-        lens_run: bool             = true,
+        lens_debug: bool            = true,
+        lens_enable: bool           = true,
+        lens_implementations: bool  = true,
+        lens_run: bool              = true,
+        lens_methodReferences: bool = false,
 
         linkedProjects: Vec<ManifestOrProjectJson> = Vec::new(),
         lruCapacity: Option<usize>                 = None,

crates/rust-analyzer/src/handlers.rs

@@ -11,6 +11,7 @@ use ide::{
     FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData, NavigationTarget, Query,
     RangeInfo, Runnable, RunnableKind, SearchScope, TextEdit,
 };
+use itertools::Itertools;
 use lsp_server::ErrorCode;
 use lsp_types::{
     CallHierarchyIncomingCall, CallHierarchyIncomingCallsParams, CallHierarchyItem,
@@ -952,13 +953,30 @@ pub(crate) fn handle_code_lens(
                 }),
         );
     }
+
+    if snap.config.lens.references() {
+        lenses.extend(snap.analysis.find_all_methods(file_id)?.into_iter().map(|it| {
+            let range = to_proto::range(&line_index, it.range);
+            let position = to_proto::position(&line_index, it.range.start());
+            let lens_params =
+                lsp_types::TextDocumentPositionParams::new(params.text_document.clone(), position);
+
+            CodeLens {
+                range,
+                command: None,
+                data: Some(to_value(CodeLensResolveData::References(lens_params)).unwrap()),
+            }
+        }));
+    }
+
     Ok(Some(lenses))
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 enum CodeLensResolveData {
     Impls(lsp_types::request::GotoImplementationParams),
+    References(lsp_types::TextDocumentPositionParams),
 }
 
 pub(crate) fn handle_code_lens_resolve(
@@ -990,6 +1008,34 @@ pub(crate) fn handle_code_lens_resolve(
             );
             Ok(CodeLens { range: code_lens.range, command: Some(cmd), data: None })
         }
+        Some(CodeLensResolveData::References(doc_position)) => {
+            let position = from_proto::file_position(&snap, doc_position.clone())?;
+            let locations = snap
+                .analysis
+                .find_all_refs(position, None)
+                .unwrap_or(None)
+                .map(|r| {
+                    r.references()
+                        .iter()
+                        .filter_map(|it| to_proto::location(&snap, it.file_range).ok())
+                        .collect_vec()
+                })
+                .unwrap_or_default();
+
+            let title = reference_title(locations.len());
+            let cmd = if locations.is_empty() {
+                Command { title, command: "".into(), arguments: None }
+            } else {
+                show_references_command(
+                    title,
+                    &doc_position.text_document.uri,
+                    code_lens.range.start,
+                    locations,
+                )
+            };
+
+            Ok(CodeLens { range: code_lens.range, command: Some(cmd), data: None })
+        }
         None => Ok(CodeLens {
             range: code_lens.range,
             command: Some(Command { title: "Error".into(), ..Default::default() }),
@@ -1248,6 +1294,14 @@ fn implementation_title(count: usize) -> String {
     }
 }
 
+fn reference_title(count: usize) -> String {
+    if count == 1 {
+        "1 reference".into()
+    } else {
+        format!("{} references", count)
+    }
+}
+
 fn show_references_command(
     title: String,
     uri: &lsp_types::Url,

editors/code/package.json

@@ -554,6 +554,11 @@
                     "type": "boolean",
                     "default": true
                 },
+                "rust-analyzer.lens.methodReferences": {
+                    "markdownDescription": "Whether to show `Method References` lens. Only applies when `#rust-analyzer.lens.enable#` is set.",
+                    "type": "boolean",
+                    "default": false
+                },
                 "rust-analyzer.hoverActions.enable": {
                     "description": "Whether to show HoverActions in Rust files.",
                     "type": "boolean",

editors/code/src/config.ts

@@ -138,6 +138,7 @@ export class Config {
             run: this.get<boolean>("lens.run"),
             debug: this.get<boolean>("lens.debug"),
             implementations: this.get<boolean>("lens.implementations"),
+            methodReferences: this.get<boolean>("lens.methodReferences"),
         };
     }