feat: lsp rename/find-all-references for traits (#5409)

# Description

## Problem

Resolves #5410

## Summary


https://github.com/noir-lang/noir/assets/209371/6ed5b6ee-8355-4374-9626-9aacea2b131a


## 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 <15848336+TomAFrench@users.noreply.github.com>

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -30,7 +30,8 @@ use crate::{
         SecondaryAttribute, StructId,
     },
     node_interner::{
-        DefinitionId, DefinitionKind, DependencyId, ExprId, FuncId, GlobalId, TraitId, TypeAliasId,
+        DefinitionId, DefinitionKind, DependencyId, ExprId, FuncId, GlobalId, ReferenceId, TraitId,
+        TypeAliasId,
     },
     parser::TopLevelStatement,
     Shared, Type, TypeBindings, TypeVariable,
@@ -1407,7 +1408,8 @@ impl<'context> Elaborator<'context> {
             self.file = trait_impl.file_id;
             self.local_module = trait_impl.module_id;
 
-            trait_impl.trait_id = self.resolve_trait_by_path(trait_impl.trait_path.clone());
+            let trait_id = self.resolve_trait_by_path(trait_impl.trait_path.clone());
+            trait_impl.trait_id = trait_id;
             let unresolved_type = &trait_impl.object_type;
 
             self.add_generics(&trait_impl.generics);
@@ -1445,6 +1447,15 @@ impl<'context> Elaborator<'context> {
             trait_impl.resolved_object_type = self.self_type.take();
             trait_impl.impl_id = self.current_trait_impl.take();
             self.generics.clear();
+
+            if let Some(trait_id) = trait_id {
+                let referenced = ReferenceId::Trait(trait_id);
+                let reference = ReferenceId::Variable(Location::new(
+                    trait_impl.trait_path.last_segment().span(),
+                    trait_impl.file_id,
+                ));
+                self.interner.add_reference(referenced, reference);
+            }
         }
     }
 

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

@@ -519,6 +519,10 @@ fn add_import_reference(
             let variable = ReferenceId::Variable(Location::new(name.span(), file_id));
             interner.add_reference(ReferenceId::Struct(struct_id), variable);
         }
+        crate::macros_api::ModuleDefId::TraitId(trait_id) => {
+            let variable = ReferenceId::Variable(Location::new(name.span(), file_id));
+            interner.add_reference(ReferenceId::Trait(trait_id), variable);
+        }
         _ => (),
     }
 }

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

@@ -381,6 +381,7 @@ impl<'a> ModCollector<'a> {
         let mut errors: Vec<(CompilationError, FileId)> = vec![];
         for trait_definition in traits {
             let name = trait_definition.name.clone();
+            let name_location = Location::new(name.span(), self.file_id);
 
             // Create the corresponding module for the trait namespace
             let trait_id = match self.push_child_module(
@@ -532,6 +533,9 @@ impl<'a> ModCollector<'a> {
             };
             context.def_interner.push_empty_trait(trait_id, &unresolved, resolved_generics);
 
+            context.def_interner.add_trait_location(trait_id, name_location);
+            context.def_interner.add_definition_location(ReferenceId::Trait(trait_id));
+
             self.def_collector.items.traits.insert(trait_id, unresolved);
         }
         errors

compiler/noirc_frontend/src/locations.rs

@@ -34,7 +34,7 @@ impl NodeInterner {
             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!(),
+            ReferenceId::Trait(id) => self.trait_location(&id),
             ReferenceId::Global(id) => self.get_global(id).location,
             ReferenceId::Alias(id) => self.get_type_alias(id).borrow().location,
             ReferenceId::Variable(location) => location,
@@ -101,11 +101,8 @@ impl NodeInterner {
 
         let reference_node = self.reference_graph[node_index];
         let found_locations: Vec<Location> = match reference_node {
-            ReferenceId::Alias(_)
-            | ReferenceId::Global(_)
-            | ReferenceId::Module(_)
-            | ReferenceId::Trait(_) => todo!(),
-            ReferenceId::Function(_) | ReferenceId::Struct(_) => {
+            ReferenceId::Alias(_) | ReferenceId::Global(_) | ReferenceId::Module(_) => todo!(),
+            ReferenceId::Function(_) | ReferenceId::Struct(_) | ReferenceId::Trait(_) => {
                 self.find_all_references_for_index(node_index, include_reference)
             }
 

compiler/noirc_frontend/src/node_interner.rs

@@ -71,6 +71,9 @@ pub struct NodeInterner {
     // The location of each struct name
     struct_name_locations: HashMap<StructId, Location>,
 
+    // The location of each trait name
+    trait_name_locations: HashMap<TraitId, Location>,
+
     /// This graph tracks dependencies between different global definitions.
     /// This is used to ensure the absence of dependency cycles for globals and types.
     dependency_graph: DiGraph<DependencyId, ()>,
@@ -545,6 +548,7 @@ impl Default for NodeInterner {
             function_modules: HashMap::new(),
             module_locations: HashMap::new(),
             struct_name_locations: HashMap::new(),
+            trait_name_locations: HashMap::new(),
             func_id_to_trait: HashMap::new(),
             dependency_graph: petgraph::graph::DiGraph::new(),
             dependency_graph_indices: HashMap::new(),
@@ -984,6 +988,14 @@ impl NodeInterner {
         self.struct_name_locations[struct_id]
     }
 
+    pub fn add_trait_location(&mut self, trait_id: TraitId, location: Location) {
+        self.trait_name_locations.insert(trait_id, location);
+    }
+
+    pub fn trait_location(&self, trait_id: &TraitId) -> Location {
+        self.trait_name_locations[trait_id]
+    }
+
     pub fn add_module_location(&mut self, module_id: ModuleId, location: Location) {
         self.module_locations.insert(module_id, location);
     }

tooling/lsp/src/requests/rename.rs

@@ -145,4 +145,9 @@ mod rename_tests {
     async fn test_rename_struct() {
         check_rename_succeeds("rename_struct", "Foo").await;
     }
+
+    #[test]
+    async fn test_rename_trait() {
+        check_rename_succeeds("rename_trait", "Foo").await;
+    }
 }

tooling/lsp/test_programs/rename_trait/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "rename_trait"
+type = "bin"
+authors = [""]
+
+[dependencies]

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

@@ -0,0 +1,21 @@
+mod foo {
+    mod bar {
+        trait Foo {
+            fn foo() {}
+        }
+    }
+}
+
+use foo::bar::Foo;
+
+struct Bar {
+
+}
+
+impl Foo for Bar {
+
+}
+
+fn main() {
+    foo::bar::Foo::foo();
+}