feat: lsp "go to definition" for modules (#5406)

# Description

## Problem

Resolves #5405

## Summary

Now clicking on module segments works as expected.


https://github.com/noir-lang/noir/assets/209371/e8968eea-ad0b-4d85-9500-db31e46009a0

## Additional Context



## 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/elaborator/scope.rs

@@ -48,11 +48,11 @@ impl<'context> Elaborator<'context> {
         let path_resolution;
 
         if self.interner.track_references {
-            let mut dependencies: Vec<ReferenceId> = Vec::new();
+            let mut references: Vec<ReferenceId> = Vec::new();
             path_resolution =
-                resolver.resolve(self.def_maps, path.clone(), &mut Some(&mut dependencies))?;
+                resolver.resolve(self.def_maps, path.clone(), &mut Some(&mut references))?;
 
-            for (referenced, ident) in dependencies.iter().zip(path.segments) {
+            for (referenced, ident) in references.iter().zip(path.segments) {
                 let reference = ReferenceId::Variable(Location::new(ident.span(), self.file));
                 self.interner.add_reference(*referenced, reference);
             }

compiler/noirc_frontend/src/hir/def_collector/dc_crate.rs

@@ -330,7 +330,28 @@ impl DefCollector {
         // Resolve unresolved imports collected from the crate, one by one.
         for collected_import in std::mem::take(&mut def_collector.imports) {
             let module_id = collected_import.module_id;
-            match resolve_import(crate_id, &collected_import, &context.def_maps, &mut None) {
+            let resolved_import = if context.def_interner.track_references {
+                let mut references: Vec<ReferenceId> = Vec::new();
+                let resolved_import = resolve_import(
+                    crate_id,
+                    &collected_import,
+                    &context.def_maps,
+                    &mut Some(&mut references),
+                );
+
+                let current_def_map = context.def_maps.get(&crate_id).unwrap();
+                let file_id = current_def_map.file_id(module_id);
+
+                for (referenced, ident) in references.iter().zip(&collected_import.path.segments) {
+                    let reference = ReferenceId::Variable(Location::new(ident.span(), file_id));
+                    context.def_interner.add_reference(*referenced, reference);
+                }
+
+                resolved_import
+            } else {
+                resolve_import(crate_id, &collected_import, &context.def_maps, &mut None)
+            };
+            match resolved_import {
                 Ok(resolved_import) => {
                     if let Some(error) = resolved_import.error {
                         errors.push((

compiler/noirc_frontend/src/hir/def_collector/dc_mod.rs

@@ -3,7 +3,7 @@ use std::vec;
 
 use acvm::{AcirField, FieldElement};
 use fm::{FileId, FileManager, FILE_EXTENSION};
-use noirc_errors::Location;
+use noirc_errors::{Location, Span};
 use num_bigint::BigUint;
 use num_traits::Num;
 use rustc_hash::FxHashMap as HashMap;
@@ -283,12 +283,18 @@ impl<'a> ModCollector<'a> {
             );
 
             // Create the corresponding module for the struct namespace
-            let id = match self.push_child_module(&name, self.file_id, false, false) {
-                Ok(local_id) => context.def_interner.new_struct(
+            let id = match self.push_child_module(
+                context,
+                &name,
+                Location::new(name.span(), self.file_id),
+                false,
+                false,
+            ) {
+                Ok(module_id) => context.def_interner.new_struct(
                     &unresolved,
                     resolved_generics,
                     krate,
-                    local_id,
+                    module_id.local_id,
                     self.file_id,
                 ),
                 Err(error) => {
@@ -377,8 +383,14 @@ impl<'a> ModCollector<'a> {
             let name = trait_definition.name.clone();
 
             // Create the corresponding module for the trait namespace
-            let trait_id = match self.push_child_module(&name, self.file_id, false, false) {
-                Ok(local_id) => TraitId(ModuleId { krate, local_id }),
+            let trait_id = match self.push_child_module(
+                context,
+                &name,
+                Location::new(name.span(), self.file_id),
+                false,
+                false,
+            ) {
+                Ok(module_id) => TraitId(ModuleId { krate, local_id: module_id.local_id }),
                 Err(error) => {
                     errors.push((error.into(), self.file_id));
                     continue;
@@ -535,13 +547,19 @@ impl<'a> ModCollector<'a> {
     ) -> Vec<(CompilationError, FileId)> {
         let mut errors: Vec<(CompilationError, FileId)> = vec![];
         for submodule in submodules {
-            match self.push_child_module(&submodule.name, file_id, true, submodule.is_contract) {
+            match self.push_child_module(
+                context,
+                &submodule.name,
+                Location::new(submodule.name.span(), file_id),
+                true,
+                submodule.is_contract,
+            ) {
                 Ok(child) => {
                     errors.extend(collect_defs(
                         self.def_collector,
                         submodule.contents,
                         file_id,
-                        child,
+                        child.local_id,
                         crate_id,
                         context,
                         macro_processors,
@@ -620,13 +638,24 @@ impl<'a> ModCollector<'a> {
         );
 
         // Add module into def collector and get a ModuleId
-        match self.push_child_module(&mod_decl.ident, child_file_id, true, false) {
+        match self.push_child_module(
+            context,
+            &mod_decl.ident,
+            Location::new(Span::empty(0), child_file_id),
+            true,
+            false,
+        ) {
             Ok(child_mod_id) => {
+                // Track that the "foo" in `mod foo;` points to the module "foo"
+                let referenced = ReferenceId::Module(child_mod_id);
+                let reference = ReferenceId::Variable(location);
+                context.def_interner.add_reference(referenced, reference);
+
                 errors.extend(collect_defs(
                     self.def_collector,
                     ast,
                     child_file_id,
-                    child_mod_id,
+                    child_mod_id.local_id,
                     crate_id,
                     context,
                     macro_processors,
@@ -643,13 +672,22 @@ impl<'a> ModCollector<'a> {
     /// On error this returns None and pushes to `errors`
     fn push_child_module(
         &mut self,
+        context: &mut Context,
         mod_name: &Ident,
-        file_id: FileId,
+        mod_location: Location,
         add_to_parent_scope: bool,
         is_contract: bool,
-    ) -> Result<LocalModuleId, DefCollectorErrorKind> {
+    ) -> Result<ModuleId, DefCollectorErrorKind> {
         let parent = Some(self.module_id);
-        let location = Location::new(mod_name.span(), file_id);
+
+        // Note: the difference between `location` and `mod_location` is:
+        // - `mod_location` will point to either the token "foo" in `mod foo { ... }`
+        //   if it's an inline module, or the first char of a the file if it's an external module.
+        // - `location` will always point to the token "foo" in `mod foo` regardless of whether
+        //   it's inline or external.
+        // Eventually the location put in `ModuleData` is used for codelenses about `contract`s,
+        // so we keep using `location` so that it continues to work as usual.
+        let location = Location::new(mod_name.span(), mod_location.file);
         let new_module = ModuleData::new(parent, location, is_contract);
         let module_id = self.def_collector.def_map.modules.insert(new_module);
 
@@ -658,6 +696,11 @@ impl<'a> ModCollector<'a> {
         // Update the parent module to reference the child
         modules[self.module_id.0].children.insert(mod_name.clone(), LocalModuleId(module_id));
 
+        let mod_id = ModuleId {
+            krate: self.def_collector.def_map.krate,
+            local_id: LocalModuleId(module_id),
+        };
+
         // Add this child module into the scope of the parent module as a module definition
         // module definitions are definitions which can only exist at the module level.
         // ModuleDefinitionIds can be used across crates since they contain the CrateId
@@ -666,11 +709,6 @@ impl<'a> ModCollector<'a> {
         // to a child module containing its methods) since the module name should not shadow
         // the struct name.
         if add_to_parent_scope {
-            let mod_id = ModuleId {
-                krate: self.def_collector.def_map.krate,
-                local_id: LocalModuleId(module_id),
-            };
-
             if let Err((first_def, second_def)) =
                 modules[self.module_id.0].declare_child_module(mod_name.to_owned(), mod_id)
             {
@@ -681,9 +719,11 @@ impl<'a> ModCollector<'a> {
                 };
                 return Err(err);
             }
+
+            context.def_interner.add_module_location(mod_id, mod_location);
         }
 
-        Ok(LocalModuleId(module_id))
+        Ok(mod_id)
     }
 }
 

compiler/noirc_frontend/src/locations.rs

@@ -31,7 +31,7 @@ impl LocationIndices {
 impl NodeInterner {
     pub fn reference_location(&self, reference: ReferenceId) -> Location {
         match reference {
-            ReferenceId::Module(_) => todo!(),
+            ReferenceId::Module(id) => self.module_location(&id),
             ReferenceId::Function(id) => self.function_modifiers(&id).name_location,
             ReferenceId::Struct(id) => self.struct_location(&id),
             ReferenceId::Trait(_) => todo!(),

compiler/noirc_frontend/src/node_interner.rs

@@ -65,6 +65,9 @@ pub struct NodeInterner {
     // Contains the source module each function was defined in
     function_modules: HashMap<FuncId, ModuleId>,
 
+    // The location of each module
+    module_locations: HashMap<ModuleId, Location>,
+
     // The location of each struct name
     struct_name_locations: HashMap<StructId, Location>,
 
@@ -540,6 +543,7 @@ impl Default for NodeInterner {
             function_definition_ids: HashMap::new(),
             function_modifiers: HashMap::new(),
             function_modules: HashMap::new(),
+            module_locations: HashMap::new(),
             struct_name_locations: HashMap::new(),
             func_id_to_trait: HashMap::new(),
             dependency_graph: petgraph::graph::DiGraph::new(),
@@ -980,6 +984,14 @@ impl NodeInterner {
         self.struct_name_locations[struct_id]
     }
 
+    pub fn add_module_location(&mut self, module_id: ModuleId, location: Location) {
+        self.module_locations.insert(module_id, location);
+    }
+
+    pub fn module_location(&self, module_id: &ModuleId) -> Location {
+        self.module_locations[module_id]
+    }
+
     pub fn global_attributes(&self, global_id: &GlobalId) -> &[SecondaryAttribute] {
         &self.global_attributes[global_id]
     }

tooling/lsp/src/requests/goto_definition.rs

@@ -51,7 +51,7 @@ mod goto_definition_tests {
 
     use super::*;
 
-    async fn expect_goto(directory: &str, name: &str, definition_index: usize) {
+    async fn expect_goto_for_all_references(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);
@@ -90,21 +90,20 @@ mod goto_definition_tests {
         }
     }
 
-    #[test]
-    async fn goto_from_function_location_to_declaration() {
-        expect_goto("go_to_definition", "another_function", 0).await;
-    }
-
-    #[test]
-    async fn goto_from_use_as() {
-        let (mut state, noir_text_document) = test_utils::init_lsp_server("go_to_definition").await;
+    async fn expect_goto(
+        directory: &str,
+        position: Position,
+        expected_file: &str,
+        expected_range: Range,
+    ) {
+        let (mut state, noir_text_document) = test_utils::init_lsp_server(directory).await;
 
         let params = GotoDefinitionParams {
             text_document_position_params: lsp_types::TextDocumentPositionParams {
                 text_document: lsp_types::TextDocumentIdentifier {
                     uri: noir_text_document.clone(),
                 },
-                position: Position { line: 7, character: 29 }, // The word after `as`
+                position,
             },
             work_done_progress_params: Default::default(),
             partial_result_params: Default::default(),
@@ -116,15 +115,85 @@ mod goto_definition_tests {
             .unwrap_or_else(|| panic!("Didn't get a goto definition response"));
 
         if let GotoDefinitionResponse::Scalar(location) = response {
-            assert_eq!(
-                location.range,
-                Range {
-                    start: Position { line: 1, character: 11 },
-                    end: Position { line: 1, character: 27 }
-                }
-            );
+            assert!(location.uri.to_string().ends_with(expected_file));
+            assert_eq!(location.range, expected_range);
         } else {
             panic!("Expected a scalar response");
         };
     }
+
+    #[test]
+    async fn goto_from_function_location_to_declaration() {
+        expect_goto_for_all_references("go_to_definition", "another_function", 0).await;
+    }
+
+    #[test]
+    async fn goto_from_use_as() {
+        expect_goto(
+            "go_to_definition",
+            Position { line: 7, character: 29 }, // The word after `as`,
+            "src/main.nr",
+            Range {
+                start: Position { line: 1, character: 11 },
+                end: Position { line: 1, character: 27 },
+            },
+        )
+        .await;
+    }
+
+    #[test]
+    async fn goto_module_from_call_path() {
+        expect_goto(
+            "go_to_definition",
+            Position { line: 17, character: 4 }, // "bar" in "bar::baz()"
+            "src/bar.nr",
+            Range {
+                start: Position { line: 0, character: 0 },
+                end: Position { line: 0, character: 0 },
+            },
+        )
+        .await;
+    }
+
+    #[test]
+    async fn goto_inline_module_from_call_path() {
+        expect_goto(
+            "go_to_definition",
+            Position { line: 18, character: 9 }, // "inline" in "bar::inline::qux()"
+            "src/bar.nr",
+            Range {
+                start: Position { line: 2, character: 4 },
+                end: Position { line: 2, character: 10 },
+            },
+        )
+        .await;
+    }
+
+    #[test]
+    async fn goto_module_from_use_path() {
+        expect_goto(
+            "go_to_definition",
+            Position { line: 6, character: 4 }, // "foo" in "use foo::another_function;"
+            "src/main.nr",
+            Range {
+                start: Position { line: 0, character: 4 },
+                end: Position { line: 0, character: 7 },
+            },
+        )
+        .await;
+    }
+
+    #[test]
+    async fn goto_module_from_mod() {
+        expect_goto(
+            "go_to_definition",
+            Position { line: 9, character: 4 }, // "bar" in "mod bar;"
+            "src/bar.nr",
+            Range {
+                start: Position { line: 0, character: 0 },
+                end: Position { line: 0, character: 0 },
+            },
+        )
+        .await;
+    }
 }

tooling/lsp/test_programs/go_to_definition/src/bar.nr

@@ -0,0 +1,5 @@
+pub fn baz() {}
+
+mod inline {
+    pub fn qux() {}
+}

tooling/lsp/test_programs/go_to_definition/src/main.nr

@@ -7,10 +7,14 @@ mod foo {
 use foo::another_function;
 use foo::another_function as aliased_function;
 
+mod bar;
+
 fn some_function() -> Field {
     1 + 2
 }
 
 fn main() {
     let _ = another_function();
+    bar::baz();
+    bar::inline::qux();
 }