2 files changed, 56 insertions, 43 deletions

diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 4763c0aac..286a6a110 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -56,7 +56,7 @@ use ide_db::{
     symbol_index::{self, FileSymbol},
     LineIndexDatabase,
 };
-use syntax::{SourceFile, TextRange, TextSize};
+use syntax::{SourceFile, SyntaxKind, TextRange, TextSize};
 
 use crate::display::ToNav;
 
@@ -369,6 +369,21 @@ impl Analysis {
         })
     }
 
+    /// Finds all methods and free functions for the file.
+    pub fn find_all_methods(&self, file_id: FileId) -> Cancelable<Vec<FileRange>> {
+        let res = self
+            .file_structure(file_id)?
+            .into_iter()
+            .filter(|it| match it.kind {
+                SyntaxKind::FN => true,
+                _ => false,
+            })
+            .filter_map(|it| Some(FileRange { file_id, range: it.navigation_range }))
+            .collect();
+
+        Ok(res)
+    }
+
     /// Returns a short text describing element at position.
     pub fn hover(
         &self,
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 06afb8148..1a0bee5e4 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -955,48 +955,18 @@ 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);
+        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))
@@ -1006,6 +976,7 @@ pub(crate) fn handle_code_lens(
 #[serde(rename_all = "camelCase")]
 enum CodeLensResolveData {
     Impls(lsp_types::request::GotoImplementationParams),
+    References(lsp_types::TextDocumentPositionParams),
 }
 
 pub(crate) fn handle_code_lens_resolve(
@@ -1037,6 +1008,33 @@ 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 cmd = if locations.is_empty() {
+                Command { title: "No references".into(), command: "".into(), arguments: None }
+            } else {
+                show_references_command(
+                    reference_title(locations.len()),
+                    &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() }),