fix: go to definition from use statement

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

@@ -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;
         }
@@ -45,32 +45,10 @@
         }
 
         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 @@
         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

@@ -182,7 +182,7 @@
     /// Stores the [Location] of a [Type] reference
     pub(crate) type_ref_locations: Vec<(Type, Location)>,
 
     /// In Noir's metaprogramming, a noir type has the type `Type`. When these are spliced
     /// into `quoted` expressions, we preserve the original type by assigning it a unique id
     /// and creating a `Token::QuotedType(id)` from this id. We cannot create a token holding
     /// the actual type since types do not implement Send or Sync.
@@ -191,7 +191,22 @@
     /// 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 @@
         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> {
@@ -172,7 +176,7 @@
     ///
     /// ### Example:
     /// ```nr
     /// trait Fieldable {
     ///     fn to_field(self) -> Field;
     ///        ^------------------------------\
     /// }                                     |    

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;
     }
 }