fix: lsp find struct reference in return locations and paths (#5404)

# Description

## Problem

Resolves #5394

## Summary

Finds struct locations in return types and inside paths.

## Additional Context

Making it work for the return type was simple: the code to register it
is now done a bit earlier so it works in all cases where a type is
referred as the last segment in a Path.

Making it work when the type is in the middle of a Path (like in
`foo::Foo::new`) was a bit trickier: the existing Path lookup only cared
about (and registered) the result of the last segment. So now there's
code to register each segment in turn... but this is only done if
running in LSP mode. So in `foo::Foo::new`, `foo` is linked to a module
and `Foo` is linked to a struct. The existing code relied on the
`DependencyId` for this, but a "module" dependency didn't exist... to
avoid introducing new variants that are only used in this reference
tracking I instead created a new enum, `ReferenceId`.

## 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.

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 @@ use crate::ast::{
     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 @@ impl<'a> ModCollector<'a> {
             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
     }

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