Add method references CodeLens

2 files changed, 68 insertions, 8 deletions

diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index fab15f860..fa8472e62 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/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: true }
     }
 }
 
 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 = true,
 
         linkedProjects: Vec<ManifestOrProjectJson> = Vec::new(),
         lruCapacity: Option<usize>                 = None,
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index f7c7a378a..06afb8148 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/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,6 +953,52 @@ pub(crate) fn handle_code_lens(
                 }),
         );
     }
+
+    if snap.config.lens.references() {
+        let ref_lenses = snap
+            .analysis
+            .file_structure(file_id)?
+            .into_iter()
+            .filter(|it| match it.kind {
+                SyntaxKind::FN => true,
+                _ => false,
+            })
+            .filter_map(|it| {
+                let position = FilePosition { file_id, offset: it.navigation_range.start() };
+                let scope = None; // all references
+
+                snap.analysis.find_all_refs(position, scope).unwrap_or(None).map(|r| {
+                    let mut lenses = Vec::new();
+                    if r.len() == 1 {
+                        // Only a declaration
+                        return lenses;
+                    }
+
+                    let uri = to_proto::url(&snap, file_id);
+                    let range = to_proto::range(&line_index, it.node_range);
+                    let position = to_proto::position(&line_index, position.offset);
+
+                    if snap.config.lens.method_refs {
+                        let all_locations: Vec<_> = r
+                            .references()
+                            .iter()
+                            .filter_map(|it| to_proto::location(&snap, it.file_range).ok())
+                            .collect();
+                        let title = reference_title(all_locations.len());
+                        let all_refs =
+                            show_references_command(title, &uri, position, all_locations);
+                        lenses.push(CodeLens { range, command: Some(all_refs), data: None });
+                    }
+
+                    lenses
+                })
+            })
+            .flatten()
+            .collect_vec();
+
+        lenses.extend(ref_lenses);
+    }
+
     Ok(Some(lenses))
 }
 
@@ -1248,6 +1295,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,