fix: lsp find struct reference in return locations and paths

compiler/noirc_frontend/src/elaborator/expressions.rs

@@ -27,7 +27,7 @@ use crate::{
         HirLiteral, HirStatement, Ident, IndexExpression, Literal, MemberAccessExpression,
         MethodCallExpression, PrefixExpression,
     },
-    node_interner::{DefinitionKind, DependencyId, ExprId, FuncId},
+    node_interner::{DefinitionKind, ExprId, FuncId, ReferenceId},
     token::Tokens,
     Kind, QuotedType, Shared, StructType, Type,
 };
@@ -432,8 +432,8 @@ impl<'context> Elaborator<'context> {
             struct_generics,
         });
 
-        let referenced = DependencyId::Struct(struct_type.borrow().id);
-        let reference = DependencyId::Variable(Location::new(span, self.file));
+        let referenced = ReferenceId::Struct(struct_type.borrow().id);
+        let reference = ReferenceId::Variable(Location::new(span, self.file));
         self.interner.add_reference(referenced, reference);
 
         (expr, Type::Struct(struct_type, generics))

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -534,7 +534,7 @@ impl<'context> Elaborator<'context> {
     fn resolve_trait_by_path(&mut self, path: Path) -> Option<TraitId> {
         let path_resolver = StandardPathResolver::new(self.module_id());
 
-        let error = match path_resolver.resolve(self.def_maps, path.clone()) {
+        let error = match path_resolver.resolve(self.def_maps, path.clone(), &mut None) {
             Ok(PathResolution { module_def_id: ModuleDefId::TraitId(trait_id), error }) => {
                 if let Some(error) = error {
                     self.push_err(error);

compiler/noirc_frontend/src/elaborator/patterns.rs

@@ -15,7 +15,7 @@ use crate::{
         stmt::HirPattern,
     },
     macros_api::{HirExpression, Ident, Path, Pattern},
-    node_interner::{DefinitionId, DefinitionKind, DependencyId, ExprId, GlobalId, TraitImplKind},
+    node_interner::{DefinitionId, DefinitionKind, ExprId, GlobalId, ReferenceId, TraitImplKind},
     Shared, StructType, Type, TypeBindings,
 };
 
@@ -199,8 +199,8 @@ impl<'context> Elaborator<'context> {
             new_definitions,
         );
 
-        let referenced = DependencyId::Struct(struct_type.borrow().id);
-        let reference = DependencyId::Variable(Location::new(name_span, self.file));
+        let referenced = ReferenceId::Struct(struct_type.borrow().id);
+        let reference = ReferenceId::Variable(Location::new(name_span, self.file));
         self.interner.add_reference(referenced, reference);
 
         HirPattern::Struct(expected_type, fields, location)
@@ -446,8 +446,8 @@ impl<'context> Elaborator<'context> {
                             self.interner.add_function_dependency(current_item, func_id);
                         }
 
-                        let variable = DependencyId::Variable(hir_ident.location);
-                        let function = DependencyId::Function(func_id);
+                        let variable = ReferenceId::Variable(hir_ident.location);
+                        let function = ReferenceId::Function(func_id);
                         self.interner.add_reference(function, variable);
                     }
                     DefinitionKind::Global(global_id) => {
@@ -458,8 +458,8 @@ impl<'context> Elaborator<'context> {
                             self.interner.add_global_dependency(current_item, global_id);
                         }
 
-                        let variable = DependencyId::Variable(hir_ident.location);
-                        let global = DependencyId::Global(global_id);
+                        let variable = ReferenceId::Variable(hir_ident.location);
+                        let global = ReferenceId::Global(global_id);
                         self.interner.add_reference(global, variable);
                     }
                     DefinitionKind::GenericType(_) => {

compiler/noirc_frontend/src/elaborator/scope.rs

@@ -1,11 +1,12 @@
-use noirc_errors::Spanned;
+use noirc_errors::{Location, Spanned};
 
 use crate::ast::ERROR_IDENT;
 use crate::hir::def_map::{LocalModuleId, ModuleId};
 use crate::hir::resolution::path_resolver::{PathResolver, StandardPathResolver};
 use crate::hir::resolution::resolver::SELF_TYPE_NAME;
 use crate::hir::scope::{Scope as GenericScope, ScopeTree as GenericScopeTree};
 use crate::macros_api::Ident;
+use crate::node_interner::ReferenceId;
 use crate::{
     hir::{
         def_map::{ModuleDefId, TryFromModuleDefId},
@@ -44,7 +45,20 @@ impl<'context> Elaborator<'context> {
 
     pub(super) fn resolve_path(&mut self, path: Path) -> Result<ModuleDefId, ResolverError> {
         let resolver = StandardPathResolver::new(self.module_id());
-        let path_resolution = resolver.resolve(self.def_maps, path)?;
+        let path_resolution;
+
+        if self.interner.track_references {
+            let mut dependencies: Vec<ReferenceId> = Vec::new();
+            path_resolution =
+                resolver.resolve(self.def_maps, path.clone(), &mut Some(&mut dependencies))?;
+
+            for (referenced, ident) in dependencies.iter().zip(path.segments) {
+                let reference = ReferenceId::Variable(Location::new(ident.span(), self.file));
+                self.interner.add_reference(*referenced, reference);
+            }
+        } else {
+            path_resolution = resolver.resolve(self.def_maps, path, &mut None)?;
+        }
 
         if let Some(error) = path_resolution.error {
             self.push_err(error);

compiler/noirc_frontend/src/elaborator/types.rs

@@ -31,7 +31,8 @@ use crate::{
         UnaryOp, UnresolvedType, UnresolvedTypeData,
     },
     node_interner::{
-        DefinitionKind, DependencyId, ExprId, GlobalId, TraitId, TraitImplKind, TraitMethodId,
+        DefinitionKind, DependencyId, ExprId, GlobalId, ReferenceId, TraitId, TraitImplKind,
+        TraitMethodId,
     },
     Generics, Kind, ResolvedGeneric, Type, TypeBinding, TypeVariable, TypeVariableKind,
 };
@@ -146,13 +147,17 @@ impl<'context> Elaborator<'context> {
             Resolved(id) => self.interner.get_quoted_type(id).clone(),
         };
 
-        if let Type::Struct(_, _) = resolved_type {
+        if let Type::Struct(ref struct_type, _) = resolved_type {
             if let Some(unresolved_span) = typ.span {
                 // Record the location of the type reference
                 self.interner.push_type_ref_location(
                     resolved_type.clone(),
                     Location::new(unresolved_span, self.file),
                 );
+
+                let referenced = ReferenceId::Struct(struct_type.borrow().id);
+                let reference = ReferenceId::Variable(Location::new(unresolved_span, self.file));
+                self.interner.add_reference(referenced, reference);
             }
         }
 
@@ -244,8 +249,6 @@ impl<'context> Elaborator<'context> {
             return Type::Alias(alias, args);
         }
 
-        let last_segment = path.last_segment();
-
         match self.lookup_struct_or_error(path) {
             Some(struct_type) => {
                 if self.resolving_ids.contains(&struct_type.borrow().id) {
@@ -283,11 +286,6 @@ impl<'context> Elaborator<'context> {
                     self.interner.add_type_dependency(current_item, dependency_id);
                 }
 
-                let referenced = DependencyId::Struct(struct_type.borrow().id);
-                let reference =
-                    DependencyId::Variable(Location::new(last_segment.span(), self.file));
-                self.interner.add_reference(referenced, reference);
-
                 Type::Struct(struct_type, args)
             }
             None => Type::Error,

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

@@ -20,7 +20,7 @@ use crate::hir::Context;
 
 use crate::macros_api::{MacroError, MacroProcessor};
 use crate::node_interner::{
-    DependencyId, FuncId, GlobalId, NodeInterner, StructId, TraitId, TraitImplId, TypeAliasId,
+    FuncId, GlobalId, NodeInterner, ReferenceId, StructId, TraitId, TraitImplId, TypeAliasId,
 };
 
 use crate::ast::{
@@ -330,7 +330,7 @@ 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) {
+            match resolve_import(crate_id, &collected_import, &context.def_maps, &mut None) {
                 Ok(resolved_import) => {
                     if let Some(error) = resolved_import.error {
                         errors.push((
@@ -491,12 +491,12 @@ 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(DependencyId::Function(func_id), variable);
+            let variable = ReferenceId::Variable(Location::new(name.span(), file_id));
+            interner.add_reference(ReferenceId::Function(func_id), variable);
         }
         crate::macros_api::ModuleDefId::TypeId(struct_id) => {
-            let variable = DependencyId::Variable(Location::new(name.span(), file_id));
-            interner.add_reference(DependencyId::Struct(struct_id), variable);
+            let variable = ReferenceId::Variable(Location::new(name.span(), file_id));
+            interner.add_reference(ReferenceId::Struct(struct_id), variable);
         }
         _ => (),
     }
@@ -524,6 +524,7 @@ fn inject_prelude(
             &context.def_maps,
             ModuleId { krate: crate_id, local_id: crate_root },
             path,
+            &mut None,
         ) {
             assert!(error.is_none(), "Tried to add private item to prelude");
             let module_id = module_def_id.as_module().expect("std::prelude should be a module");

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

@@ -14,7 +14,7 @@
     TypeImpl,
 };
 use crate::macros_api::NodeInterner;
-use crate::node_interner::DependencyId;
+use crate::node_interner::ReferenceId;
 use crate::{
     graph::CrateId,
     hir::def_collector::dc_crate::{UnresolvedStruct, UnresolvedTrait},
@@ -314,7 +314,7 @@
             self.def_collector.items.types.insert(id, unresolved);
 
             context.def_interner.add_struct_location(id, name_location);
-            context.def_interner.add_definition_location(DependencyId::Struct(id));
+            context.def_interner.add_definition_location(ReferenceId::Struct(id));
         }
         definition_errors
     }
@@ -490,7 +490,7 @@
                         }
                     }
                     TraitItem::Type { name } => {
                         // TODO(nickysn or alexvitkov): implement context.def_interner.push_empty_type_alias and get an id, instead of using TypeAliasId::dummy_id()
                         if let Err((first_def, second_def)) = self.def_collector.def_map.modules
                             [trait_id.0.local_id.0]
                             .declare_type_alias(name.clone(), TypeAliasId::dummy_id())

compiler/noirc_frontend/src/hir/resolution/import.rs

@@ -3,6 +3,7 @@ use thiserror::Error;
 
 use crate::graph::CrateId;
 use crate::hir::def_collector::dc_crate::CompilationError;
+use crate::node_interner::ReferenceId;
 use std::collections::BTreeMap;
 
 use crate::ast::{Ident, ItemVisibility, Path, PathKind};
@@ -80,13 +81,14 @@ pub fn resolve_import(
     crate_id: CrateId,
     import_directive: &ImportDirective,
     def_maps: &BTreeMap<CrateId, CrateDefMap>,
+    path_references: &mut Option<&mut Vec<ReferenceId>>,
 ) -> Result<ResolvedImport, PathResolutionError> {
     let module_scope = import_directive.module_id;
     let NamespaceResolution {
         module_id: resolved_module,
         namespace: resolved_namespace,
         mut error,
-    } = resolve_path_to_ns(import_directive, crate_id, crate_id, def_maps)?;
+    } = resolve_path_to_ns(import_directive, crate_id, crate_id, def_maps, path_references)?;
 
     let name = resolve_path_name(import_directive);
 
@@ -124,14 +126,21 @@ fn resolve_path_to_ns(
     crate_id: CrateId,
     importing_crate: CrateId,
     def_maps: &BTreeMap<CrateId, CrateDefMap>,
+    path_references: &mut Option<&mut Vec<ReferenceId>>,
 ) -> NamespaceResolutionResult {
     let import_path = &import_directive.path.segments;
     let def_map = &def_maps[&crate_id];
 
     match import_directive.path.kind {
         crate::ast::PathKind::Crate => {
             // Resolve from the root of the crate
-            resolve_path_from_crate_root(crate_id, importing_crate, import_path, def_maps)
+            resolve_path_from_crate_root(
+                crate_id,
+                importing_crate,
+                import_path,
+                def_maps,
+                path_references,
+            )
         }
         crate::ast::PathKind::Plain => {
             // There is a possibility that the import path is empty
@@ -143,6 +152,7 @@ fn resolve_path_to_ns(
                     import_path,
                     import_directive.module_id,
                     def_maps,
+                    path_references,
                 );
             }
 
@@ -151,7 +161,13 @@ fn resolve_path_to_ns(
             let first_segment = import_path.first().expect("ice: could not fetch first segment");
             if current_mod.find_name(first_segment).is_none() {
                 // Resolve externally when first segment is unresolved
-                return resolve_external_dep(def_map, import_directive, def_maps, importing_crate);
+                return resolve_external_dep(
+                    def_map,
+                    import_directive,
+                    def_maps,
+                    path_references,
+                    importing_crate,
+                );
             }
 
             resolve_name_in_module(
@@ -160,12 +176,17 @@ fn resolve_path_to_ns(
                 import_path,
                 import_directive.module_id,
                 def_maps,
+                path_references,
             )
         }
 
-        crate::ast::PathKind::Dep => {
-            resolve_external_dep(def_map, import_directive, def_maps, importing_crate)
-        }
+        crate::ast::PathKind::Dep => resolve_external_dep(
+            def_map,
+            import_directive,
+            def_maps,
+            path_references,
+            importing_crate,
+        ),
     }
 }
 
@@ -175,13 +196,15 @@ fn resolve_path_from_crate_root(
 
     import_path: &[Ident],
     def_maps: &BTreeMap<CrateId, CrateDefMap>,
+    path_references: &mut Option<&mut Vec<ReferenceId>>,
 ) -> NamespaceResolutionResult {
     resolve_name_in_module(
         crate_id,
         importing_crate,
         import_path,
         def_maps[&crate_id].root,
         def_maps,
+        path_references,
     )
 }
 
@@ -191,6 +214,7 @@ fn resolve_name_in_module(
     import_path: &[Ident],
     starting_mod: LocalModuleId,
     def_maps: &BTreeMap<CrateId, CrateDefMap>,
+    path_references: &mut Option<&mut Vec<ReferenceId>>,
 ) -> NamespaceResolutionResult {
     let def_map = &def_maps[&krate];
     let mut current_mod_id = ModuleId { krate, local_id: starting_mod };
@@ -221,12 +245,27 @@ fn resolve_name_in_module(
 
         // In the type namespace, only Mod can be used in a path.
         current_mod_id = match typ {
-            ModuleDefId::ModuleId(id) => id,
+            ModuleDefId::ModuleId(id) => {
+                if let Some(path_references) = path_references {
+                    path_references.push(ReferenceId::Module(id));
+                }
+                id
+            }
             ModuleDefId::FunctionId(_) => panic!("functions cannot be in the type namespace"),
             // TODO: If impls are ever implemented, types can be used in a path
-            ModuleDefId::TypeId(id) => id.module_id(),
+            ModuleDefId::TypeId(id) => {
+                if let Some(path_references) = path_references {
+                    path_references.push(ReferenceId::Struct(id));
+                }
+                id.module_id()
+            }
             ModuleDefId::TypeAliasId(_) => panic!("type aliases cannot be used in type namespace"),
-            ModuleDefId::TraitId(id) => id.0,
+            ModuleDefId::TraitId(id) => {
+                if let Some(path_references) = path_references {
+                    path_references.push(ReferenceId::Trait(id));
+                }
+                id.0
+            }
             ModuleDefId::GlobalId(_) => panic!("globals cannot be in the type namespace"),
         };
 
@@ -270,6 +309,7 @@ fn resolve_external_dep(
     current_def_map: &CrateDefMap,
     directive: &ImportDirective,
     def_maps: &BTreeMap<CrateId, CrateDefMap>,
+    path_references: &mut Option<&mut Vec<ReferenceId>>,
     importing_crate: CrateId,
 ) -> NamespaceResolutionResult {
     // Use extern_prelude to get the dep
@@ -299,7 +339,7 @@ fn resolve_external_dep(
         is_prelude: false,
     };
 
-    resolve_path_to_ns(&dep_directive, dep_module.krate, importing_crate, def_maps)
+    resolve_path_to_ns(&dep_directive, dep_module.krate, importing_crate, def_maps, path_references)
 }
 
 // Issue an error if the given private function is being called from a non-child module, or