feat: lsp "find all references"

compiler/noirc_frontend/src/locations.rs

@@ -13,7 +13,7 @@
 
 impl LocationIndices {
     pub(crate) fn add_location(&mut self, location: Location, node_index: PetGraphIndex) {
         // Some location spans are empty: maybe they are from ficticious nodes?
         if location.span.start() == location.span.end() {
             return;
         }
@@ -88,31 +88,42 @@
 
     // Starting at the given location, find the node referenced by it. Then, gather
     // all locations that reference that node, and return all of them
-    // (the referenced node and the references).
+    // (the references and optionally the reference node if `include_reference` is true).
     // Returns `None` if the location is not known to this interner.
-    pub fn find_all_references(&self, location: Location) -> Option<Vec<Location>> {
+    pub fn find_all_references(
+        &self,
+        location: Location,
+        include_reference: bool,
+    ) -> Option<Vec<Location>> {
         let node_index = self.location_indices.get_node_from_location(location)?;
 
         let reference_node = self.reference_graph[node_index];
         let found_locations: Vec<Location> = match reference_node {
             DependencyId::Alias(_) | DependencyId::Global(_) => todo!(),
             DependencyId::Function(_) | DependencyId::Struct(_) => {
-                self.find_all_references_for_index(node_index)
+                self.find_all_references_for_index(node_index, include_reference)
             }
 
             DependencyId::Variable(_) => {
                 let referenced_node_index = self.referenced_index(node_index)?;
-                self.find_all_references_for_index(referenced_node_index)
+                self.find_all_references_for_index(referenced_node_index, include_reference)
             }
         };
         Some(found_locations)
     }
 
-    // Given a referenced node index, find all references to it and return their locations, together
-    // with the reference node's location.
-    fn find_all_references_for_index(&self, referenced_node_index: PetGraphIndex) -> Vec<Location> {
+    // Given a referenced node index, find all references to it and return their locations, optionally together
+    // with the reference node's location if `include_reference` is true.
+    fn find_all_references_for_index(
+        &self,
+        referenced_node_index: PetGraphIndex,
+        include_reference: bool,
+    ) -> Vec<Location> {
         let id = self.reference_graph[referenced_node_index];
-        let mut edit_locations = vec![self.dependency_location(id)];
+        let mut edit_locations = Vec::new();
+        if include_reference {
+            edit_locations.push(self.dependency_location(id));
+        }
 
         self.reference_graph
             .neighbors_directed(referenced_node_index, petgraph::Direction::Incoming)

tooling/lsp/src/lib.rs

@@ -21,7 +21,7 @@ use async_lsp::{
 use fm::{codespan_files as files, FileManager};
 use fxhash::FxHashSet;
 use lsp_types::{
-    request::{PrepareRenameRequest, Rename},
+    request::{PrepareRenameRequest, References, Rename},
     CodeLens,
 };
 use nargo::{
@@ -47,7 +47,8 @@ use notifications::{
 use requests::{
     on_code_lens_request, on_formatting, on_goto_declaration_request, on_goto_definition_request,
     on_goto_type_definition_request, on_initialize, on_prepare_rename_request,
-    on_profile_run_request, on_rename_request, on_shutdown, on_test_run_request, on_tests_request,
+    on_profile_run_request, on_references_request, on_rename_request, on_shutdown,
+    on_test_run_request, on_tests_request,
 };
 use serde_json::Value as JsonValue;
 use thiserror::Error;
@@ -125,6 +126,7 @@ impl NargoLspService {
             .request::<request::GotoDefinition, _>(on_goto_definition_request)
             .request::<request::GotoDeclaration, _>(on_goto_declaration_request)
             .request::<request::GotoTypeDefinition, _>(on_goto_type_definition_request)
+            .request::<References, _>(on_references_request)
             .request::<PrepareRenameRequest, _>(on_prepare_rename_request)
             .request::<Rename, _>(on_rename_request)
             .notification::<notification::Initialized>(on_initialized)

tooling/lsp/src/requests/mod.rs

@@ -36,6 +36,7 @@ mod code_lens_request;
 mod goto_declaration;
 mod goto_definition;
 mod profile_run;
+mod references;
 mod rename;
 mod test_run;
 mod tests;
@@ -44,8 +45,8 @@ pub(crate) use {
     code_lens_request::collect_lenses_for_package, code_lens_request::on_code_lens_request,
     goto_declaration::on_goto_declaration_request, goto_definition::on_goto_definition_request,
     goto_definition::on_goto_type_definition_request, profile_run::on_profile_run_request,
-    rename::on_prepare_rename_request, rename::on_rename_request, test_run::on_test_run_request,
-    tests::on_tests_request,
+    references::on_references_request, rename::on_prepare_rename_request,
+    rename::on_rename_request, test_run::on_test_run_request, tests::on_tests_request,
 };
 
 /// LSP client will send initialization request after the server has started.
@@ -121,6 +122,11 @@ pub(crate) fn on_initialize(
                         work_done_progress: None,
                     },
                 })),
+                references_provider: Some(lsp_types::OneOf::Right(lsp_types::ReferencesOptions {
+                    work_done_progress_options: WorkDoneProgressOptions {
+                        work_done_progress: None,
+                    },
+                })),
             },
             server_info: None,
         })

tooling/lsp/src/requests/references.rs

@@ -0,0 +1,94 @@
+use std::future::{self, Future};
+
+use async_lsp::ResponseError;
+use lsp_types::{Location, ReferenceParams};
+
+use crate::LspState;
+
+use super::{process_request, to_lsp_location};
+
+pub(crate) fn on_references_request(
+    state: &mut LspState,
+    params: ReferenceParams,
+) -> impl Future<Output = Result<Option<Vec<Location>>, ResponseError>> {
+    let result =
+        process_request(state, params.text_document_position, |location, interner, files| {
+            interner.find_all_references(location, params.context.include_declaration).map(
+                |locations| {
+                    locations
+                        .iter()
+                        .filter_map(|location| to_lsp_location(files, location.file, location.span))
+                        .collect()
+                },
+            )
+        });
+    future::ready(result)
+}
+
+#[cfg(test)]
+mod references_tests {
+    use super::*;
+    use crate::test_utils::{self, search_in_file};
+    use lsp_types::{
+        PartialResultParams, ReferenceContext, TextDocumentPositionParams, WorkDoneProgressParams,
+    };
+    use tokio::test;
+
+    async fn check_references_succeeds(
+        directory: &str,
+        name: &str,
+        declaration_index: usize,
+        include_declaration: bool,
+    ) {
+        let (mut state, noir_text_document) = test_utils::init_lsp_server(directory).await;
+
+        // First we find out all of the occurrences of `name` in the main.nr file.
+        // Note that this only works if that name doesn't show up in other places where we don't
+        // expect a rename, but we craft our tests to avoid that.
+        let ranges = search_in_file(noir_text_document.path(), name);
+
+        // Test renaming works on any instance of the symbol.
+        for target_range in &ranges {
+            let target_position = target_range.start;
+
+            let params = ReferenceParams {
+                text_document_position: TextDocumentPositionParams {
+                    text_document: lsp_types::TextDocumentIdentifier {
+                        uri: noir_text_document.clone(),
+                    },
+                    position: target_position,
+                },
+                work_done_progress_params: WorkDoneProgressParams { work_done_token: None },
+                partial_result_params: PartialResultParams { partial_result_token: None },
+                context: ReferenceContext { include_declaration },
+            };
+
+            let locations = on_references_request(&mut state, params)
+                .await
+                .expect("Could not execute on_references_request")
+                .unwrap();
+
+            let mut references_ranges: Vec<_> =
+                locations.iter().map(|location| location.range).collect();
+            references_ranges.sort_by_key(|range| range.start.line);
+
+            if include_declaration {
+                assert_eq!(ranges, references_ranges);
+            } else {
+                let mut ranges_without_declaration = ranges.clone();
+                ranges_without_declaration.remove(declaration_index);
+                assert_eq!(ranges_without_declaration, references_ranges);
+            }
+        }
+    }
+
+    #[test]
+    async fn test_on_references_request_including_declaration() {
+        check_references_succeeds("rename_function", "another_function", 0, true).await;
+    }
+
+    #[test]
+    async fn test_on_references_request_without_including_declaration() {
+        check_references_succeeds("rename_function", "another_function", 0, false).await;
+    }
+}

tooling/lsp/src/requests/rename.rs

@@ -29,7 +29,7 @@ pub(crate) fn on_rename_request(
 ) -> impl Future<Output = Result<Option<WorkspaceEdit>, ResponseError>> {
     let result =
         process_request(state, params.text_document_position, |location, interner, files| {
-            let rename_changes = interner.find_all_references(location).map(|locations| {
+            let rename_changes = interner.find_all_references(location, true).map(|locations| {
                 let rs = locations.iter().fold(
                     HashMap::new(),
                     |mut acc: HashMap<Url, Vec<TextEdit>>, location| {

tooling/lsp/src/types.rs

@@ -1,6 +1,6 @@
 use fm::FileId;
 use lsp_types::{
-    DeclarationCapability, DefinitionOptions, OneOf, RenameOptions,
+    DeclarationCapability, DefinitionOptions, OneOf, ReferencesOptions, RenameOptions,
     TypeDefinitionProviderCapability,
 };
 use noirc_driver::DebugFile;
@@ -140,6 +140,10 @@ pub(crate) struct ServerCapabilities {
     /// The server provides rename support.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub(crate) rename_provider: Option<OneOf<bool, RenameOptions>>,
+
+    /// The server provides references support.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub(crate) references_provider: Option<OneOf<bool, ReferencesOptions>>,
 }
 
 #[derive(Debug, PartialEq, Clone, Default, Deserialize, Serialize)]