feat: lsp "find all references" (#5395)

# Description

## Problem

Resolves #5393

## Summary

Implements "Find all references". The logic for it is the same as the
one for "Rename symbol" so it was straight-forward to do.


https://github.com/noir-lang/noir/assets/209371/5eef03bf-6170-48f3-beaa-87dcc86532d2

## Additional Context

In the video above, "Find all references" misses a few locations. I
captured that as a separate issue ( #5394 ) because since both features
rely on the same logic, fixing that would make this one work too... so
this PR is mostly about enabling "Find all references" rather than the
logic to actually find all references.

## Documentation

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist

- [x] I have tested the changes locally.
- [ ] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: Tom French <[email protected]>

compiler/noirc_frontend/src/locations.rs

@@ -88,31 +88,42 @@ impl NodeInterner {
 
     // 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 getting references 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)]