fix: go to definition from use statement (#5390)

# Description

## Problem

Resolves #5389

## Summary

"Go to definition" works by scanning some `id_to_location` hash map
where all expressions and statements locations are stored. However,
"use" statement locations aren't captured there.

Then, we recently implemented rename functionality which works by
building a graph that connects references to referenced nodes. This same
graph can be used for "go to definition"! So that's what this PR does.
However, it will first check the graph, then try the other approach
because I think that other approach works for trait impls and other
things which aren't captured in the reference graph. I'm still trying to
understand that part and to avoid breaking things I left it as is (there
are also no tests for those things so it's hard to know if we break
something... but I plan to add more tests in the future).


https://github.com/noir-lang/noir/assets/209371/1dd2cb7e-7f41-4065-8761-4a0478937166

## Additional Context

None.

## 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.
- [x] 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/hir/def_collector/dc_crate.rs

@@ -487,11 +487,11 @@ fn add_import_reference(
     match def_id {
         crate::macros_api::ModuleDefId::FunctionId(func_id) => {
             let variable = DependencyId::Variable(Location::new(name.span(), file_id));
-            interner.add_reference_for(DependencyId::Function(func_id), variable);
+            interner.add_reference(DependencyId::Function(func_id), variable);
         }
         crate::macros_api::ModuleDefId::TypeId(struct_id) => {
             let variable = DependencyId::Variable(Location::new(name.span(), file_id));
-            interner.add_reference_for(DependencyId::Struct(struct_id), variable);
+            interner.add_reference(DependencyId::Struct(struct_id), variable);
         }
         _ => (),
     }

compiler/noirc_frontend/src/locations.rs

@@ -45,32 +45,10 @@ impl NodeInterner {
         }
 
         let referenced_index = self.get_or_insert_reference(referenced);
-        let reference_index = self.reference_graph.add_node(reference);
-
-        let referenced_location = self.dependency_location(referenced);
-        let reference_location = self.dependency_location(reference);
-
-        self.reference_graph.add_edge(referenced_index, reference_index, ());
-        self.location_indices.add_location(referenced_location, referenced_index);
-        self.location_indices.add_location(reference_location, reference_index);
-    }
-
-    pub(crate) fn add_reference_for(
-        &mut self,
-        referenced_id: DependencyId,
-        reference: DependencyId,
-    ) {
-        if !self.track_references {
-            return;
-        }
-
-        let Some(referenced_index) = self.reference_graph_indices.get(&referenced_id) else {
-            panic!("Compiler Error: Referenced index not found")
-        };
-
         let reference_location = self.dependency_location(reference);
         let reference_index = self.reference_graph.add_node(reference);
-        self.reference_graph.add_edge(*referenced_index, reference_index, ());
+
+        self.reference_graph.add_edge(reference_index, referenced_index, ());
         self.location_indices.add_location(reference_location, reference_index);
     }
 
@@ -95,42 +73,60 @@ impl NodeInterner {
         index
     }
 
-    pub fn check_rename_possible(&self, location: Location) -> bool {
+    // Given a reference location, find the location of the referenced node.
+    pub fn find_referenced_location(&self, reference_location: Location) -> Option<Location> {
+        self.location_indices
+            .get_node_from_location(reference_location)
+            .and_then(|node_index| self.referenced_index(node_index))
+            .map(|node_index| self.dependency_location(self.reference_graph[node_index]))
+    }
+
+    // Is the given location known to this interner?
+    pub fn is_location_known(&self, location: Location) -> bool {
         self.location_indices.get_node_from_location(location).is_some()
     }
 
-    pub fn find_rename_symbols_at(&self, location: Location) -> Option<Vec<Location>> {
+    // 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).
+    // Returns `None` if the location is not known to this interner.
+    pub fn find_all_references(&self, location: Location) -> 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.get_edit_locations(node_index)
+                self.find_all_references_for_index(node_index)
             }
 
             DependencyId::Variable(_) => {
-                let referenced_node_index = self
-                    .reference_graph
-                    .neighbors_directed(node_index, petgraph::Direction::Incoming)
-                    .next()?;
-
-                self.get_edit_locations(referenced_node_index)
+                let referenced_node_index = self.referenced_index(node_index)?;
+                self.find_all_references_for_index(referenced_node_index)
             }
         };
         Some(found_locations)
     }
 
-    fn get_edit_locations(&self, referenced_node_index: PetGraphIndex) -> Vec<Location> {
+    // 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> {
         let id = self.reference_graph[referenced_node_index];
         let mut edit_locations = vec![self.dependency_location(id)];
 
         self.reference_graph
-            .neighbors_directed(referenced_node_index, petgraph::Direction::Outgoing)
+            .neighbors_directed(referenced_node_index, petgraph::Direction::Incoming)
             .for_each(|reference_node_index| {
                 let id = self.reference_graph[reference_node_index];
                 edit_locations.push(self.dependency_location(id));
             });
         edit_locations
     }
+
+    // Given a reference index, returns the referenced index, if any.
+    fn referenced_index(&self, reference_index: PetGraphIndex) -> Option<PetGraphIndex> {
+        self.reference_graph
+            .neighbors_directed(reference_index, petgraph::Direction::Outgoing)
+            .next()
+    }
 }

compiler/noirc_frontend/src/node_interner.rs

@@ -191,7 +191,22 @@ pub struct NodeInterner {
     /// Whether to track references. In regular compilations this is false, but tools set it to true.
     pub(crate) track_references: bool,
 
-    /// Store the location of the references in the graph
+    /// Store the location of the references in the graph.
+    /// Edges are directed from reference nodes to referenced nodes.
+    /// For example:
+    ///
+    /// ```
+    /// let foo = 3;
+    /// //  referenced
+    /// //   ^
+    /// //   |
+    /// //   +------------+
+    /// let bar = foo;    |
+    /// //      reference |
+    /// //         v      |
+    /// //         |      |
+    /// //         +------+
+    /// ```
     pub(crate) reference_graph: DiGraph<DependencyId, ()>,
 
     /// Tracks the index of the references in the graph

compiler/noirc_frontend/src/resolve_locations.rs

@@ -38,12 +38,16 @@ impl NodeInterner {
         location: Location,
         return_type_location_instead: bool,
     ) -> Option<Location> {
-        self.find_location_index(location)
-            .and_then(|index| self.resolve_location(index, return_type_location_instead))
-            .or_else(|| self.try_resolve_trait_impl_location(location))
-            .or_else(|| self.try_resolve_trait_method_declaration(location))
-            .or_else(|| self.try_resolve_type_ref(location))
-            .or_else(|| self.try_resolve_type_alias(location))
+        // First try to find the location in the reference graph
+        self.find_referenced_location(location).or_else(|| {
+            // Otherwise fallback to the location indices
+            self.find_location_index(location)
+                .and_then(|index| self.resolve_location(index, return_type_location_instead))
+                .or_else(|| self.try_resolve_trait_impl_location(location))
+                .or_else(|| self.try_resolve_trait_method_declaration(location))
+                .or_else(|| self.try_resolve_type_ref(location))
+                .or_else(|| self.try_resolve_type_alias(location))
+        })
     }
 
     pub fn get_declaration_location_from(&self, location: Location) -> Option<Location> {

tooling/lsp/src/requests/goto_declaration.rs

@@ -2,16 +2,11 @@ use std::future::{self, Future};
 
 use crate::types::GotoDeclarationResult;
 use crate::LspState;
-use crate::{parse_diff, resolve_workspace_for_source_path};
-use async_lsp::{ErrorCode, ResponseError};
+use async_lsp::ResponseError;
 
-use fm::PathString;
 use lsp_types::request::{GotoDeclarationParams, GotoDeclarationResponse};
 
-use nargo::insert_all_files_for_workspace_into_file_manager;
-use noirc_driver::file_manager_with_stdlib;
-
-use super::{position_to_location, to_lsp_location};
+use super::{process_request, to_lsp_location};
 
 pub(crate) fn on_goto_declaration_request(
     state: &mut LspState,
@@ -25,42 +20,12 @@ fn on_goto_definition_inner(
     state: &mut LspState,
     params: GotoDeclarationParams,
 ) -> Result<GotoDeclarationResult, ResponseError> {
-    let file_path =
-        params.text_document_position_params.text_document.uri.to_file_path().map_err(|_| {
-            ResponseError::new(ErrorCode::REQUEST_FAILED, "URI is not a valid file path")
-        })?;
-
-    let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
-    let package = workspace.members.first().unwrap();
-
-    let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
-    insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_diff(&workspace_file_manager, state);
-
-    let (mut context, crate_id) =
-        crate::prepare_package(&workspace_file_manager, &parsed_files, package);
-
-    let package_root_path = package.root_dir.as_os_str().to_string_lossy().into_owned();
-    let interner = if let Some(def_interner) = state.cached_definitions.get(&package_root_path) {
-        def_interner
-    } else {
-        // We ignore the warnings and errors produced by compilation while resolving the definition
-        let _ = noirc_driver::check_crate(&mut context, crate_id, false, false, false);
-        &context.def_interner
-    };
-
-    let files = workspace_file_manager.as_file_map();
-    let file_path = PathString::from(file_path);
-    let search_for_location =
-        position_to_location(files, &file_path, &params.text_document_position_params.position)?;
-
-    let goto_declaration_response =
-        interner.get_declaration_location_from(search_for_location).and_then(|found_location| {
+    process_request(state, params.text_document_position_params, |location, interner, files| {
+        interner.get_declaration_location_from(location).and_then(|found_location| {
             let file_id = found_location.file;
             let definition_position = to_lsp_location(files, file_id, found_location.span)?;
             let response = GotoDeclarationResponse::from(definition_position).to_owned();
             Some(response)
-        });
-
-    Ok(goto_declaration_response)
+        })
+    })
 }

tooling/lsp/src/requests/goto_definition.rs

@@ -1,16 +1,12 @@
 use std::future::{self, Future};
 
-use crate::{parse_diff, resolve_workspace_for_source_path};
 use crate::{types::GotoDefinitionResult, LspState};
-use async_lsp::{ErrorCode, ResponseError};
+use async_lsp::ResponseError;
 
-use fm::PathString;
 use lsp_types::request::GotoTypeDefinitionParams;
 use lsp_types::{GotoDefinitionParams, GotoDefinitionResponse};
-use nargo::insert_all_files_for_workspace_into_file_manager;
-use noirc_driver::file_manager_with_stdlib;
 
-use super::{position_to_location, to_lsp_location};
+use super::{process_request, to_lsp_location};
 
 pub(crate) fn on_goto_definition_request(
     state: &mut LspState,
@@ -33,85 +29,68 @@ fn on_goto_definition_inner(
     params: GotoDefinitionParams,
     return_type_location_instead: bool,
 ) -> Result<GotoDefinitionResult, ResponseError> {
-    let file_path =
-        params.text_document_position_params.text_document.uri.to_file_path().map_err(|_| {
-            ResponseError::new(ErrorCode::REQUEST_FAILED, "URI is not a valid file path")
-        })?;
-
-    let workspace = resolve_workspace_for_source_path(file_path.as_path()).unwrap();
-    let package = workspace.members.first().unwrap();
-
-    let mut workspace_file_manager = file_manager_with_stdlib(&workspace.root_dir);
-    insert_all_files_for_workspace_into_file_manager(&workspace, &mut workspace_file_manager);
-    let parsed_files = parse_diff(&workspace_file_manager, state);
-
-    let (mut context, crate_id) =
-        crate::prepare_package(&workspace_file_manager, &parsed_files, package);
-
-    let package_root_path = package.root_dir.as_os_str().to_string_lossy().into_owned();
-    let interner = if let Some(def_interner) = state.cached_definitions.get(&package_root_path) {
-        def_interner
-    } else {
-        // We ignore the warnings and errors produced by compilation while resolving the definition
-        let _ = noirc_driver::check_crate(&mut context, crate_id, false, false, false);
-        &context.def_interner
-    };
-
-    let files = workspace_file_manager.as_file_map();
-    let file_path = PathString::from(file_path);
-    let search_for_location =
-        position_to_location(files, &file_path, &params.text_document_position_params.position)?;
-
-    let goto_definition_response = interner
-        .get_definition_location_from(search_for_location, return_type_location_instead)
-        .and_then(|found_location| {
-            let file_id = found_location.file;
-            let definition_position = to_lsp_location(files, file_id, found_location.span)?;
-            let response = GotoDefinitionResponse::from(definition_position).to_owned();
-            Some(response)
-        });
-
-    Ok(goto_definition_response)
+    process_request(state, params.text_document_position_params, |location, interner, files| {
+        interner.get_definition_location_from(location, return_type_location_instead).and_then(
+            |found_location| {
+                let file_id = found_location.file;
+                let definition_position = to_lsp_location(files, file_id, found_location.span)?;
+                let response = GotoDefinitionResponse::from(definition_position).to_owned();
+                Some(response)
+            },
+        )
+    })
 }
 
 #[cfg(test)]
 mod goto_definition_tests {
     use std::panic;
 
-    use crate::test_utils;
-    use lsp_types::{Position, Range};
+    use crate::test_utils::{self, search_in_file};
     use tokio::test;
 
     use super::*;
 
-    #[test]
-    async fn test_on_goto_definition() {
-        let (mut state, noir_text_document) = test_utils::init_lsp_server("go_to_definition").await;
+    async fn expect_goto(directory: &str, name: &str, definition_index: usize) {
+        let (mut state, noir_text_document) = test_utils::init_lsp_server(directory).await;
+
+        let ranges = search_in_file(noir_text_document.path(), name);
+        let expected_range = ranges[definition_index];
+
+        for (index, range) in ranges.iter().enumerate() {
+            // Ideally "go to" at the definition should return the same location, but this isn't currently
+            // working. But it's also not that important, so we'll keep it for later.
+            if index == definition_index {
+                continue;
+            }
+
+            let params = GotoDefinitionParams {
+                text_document_position_params: lsp_types::TextDocumentPositionParams {
+                    text_document: lsp_types::TextDocumentIdentifier {
+                        uri: noir_text_document.clone(),
+                    },
+                    position: range.start,
+                },
+                work_done_progress_params: Default::default(),
+                partial_result_params: Default::default(),
+            };
+
+            let response = on_goto_definition_request(&mut state, params)
+                .await
+                .expect("Could execute on_goto_definition_request")
+                .unwrap_or_else(|| {
+                    panic!("Didn't get a goto definition response for index {index}")
+                });
+
+            if let GotoDefinitionResponse::Scalar(location) = response {
+                assert_eq!(location.range, expected_range);
+            } else {
+                panic!("Expected a scalar response");
+            };
+        }
+    }
 
-        let params = GotoDefinitionParams {
-            text_document_position_params: lsp_types::TextDocumentPositionParams {
-                text_document: lsp_types::TextDocumentIdentifier { uri: noir_text_document },
-                position: Position { line: 9, character: 12 }, // Right at the beginning of "another_function"
-            },
-            work_done_progress_params: Default::default(),
-            partial_result_params: Default::default(),
-        };
-
-        let response: GotoDefinitionResponse = on_goto_definition_request(&mut state, params)
-            .await
-            .expect("Could execute on_goto_definition_request")
-            .expect("Didn't get a goto definition response");
-
-        if let GotoDefinitionResponse::Scalar(location) = response {
-            assert_eq!(
-                location.range,
-                Range {
-                    start: Position { line: 4, character: 3 },
-                    end: Position { line: 4, character: 19 },
-                }
-            );
-        } else {
-            panic!("Expected a scalar response");
-        };
+    #[test]
+    async fn goto_from_function_location_to_declaration() {
+        expect_goto("go_to_definition", "another_function", 0).await;
     }
 }