chore(LSP): several improvements for enums and enum variants

compiler/noirc_frontend/src/elaborator/enums.rs

@@ -70,6 +70,7 @@ impl Elaborator<'_> {
             type_id: Some(type_id),
             trait_id: None,
             trait_impl: None,
+            enum_variant_index: Some(variant_index),
             is_entry_point: false,
             has_inline_attribute: false,
             function_body: FunctionBody::Resolved,

compiler/noirc_frontend/src/elaborator/expressions.rs

@@ -666,7 +666,7 @@ impl<'context> Elaborator<'context> {
 
         let struct_id = struct_type.borrow().id;
         let reference_location = Location::new(last_segment.ident.span(), self.file);
-        self.interner.add_struct_reference(struct_id, reference_location, is_self_type);
+        self.interner.add_type_reference(struct_id, reference_location, is_self_type);
 
         (expr, Type::DataType(struct_type, generics))
     }

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -11,26 +11,24 @@ use crate::{
     },
     graph::CrateId,
     hir::{
-        def_collector::dc_crate::{
-            filter_literal_globals, CompilationError, ImplMap, UnresolvedFunctions,
-            UnresolvedGlobal, UnresolvedStruct, UnresolvedTraitImpl, UnresolvedTypeAlias,
-        },
         def_collector::{
-            dc_crate::{CollectedItems, UnresolvedEnum},
+            dc_crate::{
+                filter_literal_globals, CollectedItems, CompilationError, ImplMap, UnresolvedEnum,
+                UnresolvedFunctions, UnresolvedGlobal, UnresolvedStruct, UnresolvedTraitImpl,
+                UnresolvedTypeAlias,
+            },
             errors::DefCollectorErrorKind,
         },
-        def_map::{DefMaps, ModuleData},
-        def_map::{LocalModuleId, ModuleId, MAIN_FUNCTION},
+        def_map::{DefMaps, LocalModuleId, ModuleData, ModuleId, MAIN_FUNCTION},
         resolution::errors::ResolverError,
         scope::ScopeForest as GenericScopeForest,
         type_check::{generics::TraitGenerics, TypeCheckError},
         Context,
     },
-    hir_def::traits::TraitImpl,
     hir_def::{
         expr::{HirCapturedVar, HirIdent},
         function::{FuncMeta, FunctionBody, HirFunction},
-        traits::TraitConstraint,
+        traits::{TraitConstraint, TraitImpl},
         types::{Generics, Kind, ResolvedGeneric},
     },
     node_interner::{
@@ -1006,6 +1004,7 @@ impl<'context> Elaborator<'context> {
             type_id: struct_id,
             trait_id,
             trait_impl: self.current_trait_impl,
+            enum_variant_index: None,
             parameters: parameters.into(),
             parameter_idents,
             return_type: func.def.return_type.clone(),
@@ -1839,9 +1838,11 @@ impl<'context> Elaborator<'context> {
                 span: typ.enum_def.span,
             };
 
-            for (i, variant) in typ.enum_def.variants.iter().enumerate() {
-                self.interner.add_definition_location(ReferenceId::EnumVariant(*type_id, i), None);
+            datatype.borrow_mut().init_variants();
+
+            let module_id = ModuleId { krate: self.crate_id, local_id: typ.module_id };
 
+            for (i, variant) in typ.enum_def.variants.iter().enumerate() {
                 let types = vecmap(&variant.item.parameters, |typ| self.resolve_type(typ.clone()));
                 let name = variant.item.name.clone();
                 datatype.borrow_mut().push_variant(EnumVariant::new(name, types.clone()));
@@ -1857,6 +1858,9 @@ impl<'context> Elaborator<'context> {
                     &self_type,
                     unresolved.clone(),
                 );
+
+                let reference_id = ReferenceId::EnumVariant(*type_id, i);
+                self.interner.add_definition_location(reference_id, Some(module_id));
             }
         }
     }

compiler/noirc_frontend/src/elaborator/patterns.rs

@@ -238,7 +238,7 @@ impl<'context> Elaborator<'context> {
         let struct_id = struct_type.borrow().id;
 
         let reference_location = Location::new(name_span, self.file);
-        self.interner.add_struct_reference(struct_id, reference_location, is_self_type);
+        self.interner.add_type_reference(struct_id, reference_location, is_self_type);
 
         for (field_index, field) in fields.iter().enumerate() {
             let reference_location = Location::new(field.0.span(), self.file);

compiler/noirc_frontend/src/elaborator/types.rs

@@ -158,7 +158,7 @@ impl<'context> Elaborator<'context> {
                 // Record the location of the type reference
                 self.interner.push_type_ref_location(resolved_type.clone(), location);
                 if !is_synthetic {
-                    self.interner.add_struct_reference(
+                    self.interner.add_type_reference(
                         data_type.borrow().id,
                         location,
                         is_self_type_name,

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

@@ -336,12 +336,12 @@ impl<'a> ModCollector<'a> {
         krate: CrateId,
     ) -> Vec<(CompilationError, FileId)> {
         let mut definition_errors = vec![];
-        for struct_definition in types {
+        for enum_definition in types {
             if let Some((id, the_enum)) = collect_enum(
                 &mut context.def_interner,
                 &mut self.def_collector.def_map,
                 &mut context.usage_tracker,
-                struct_definition,
+                enum_definition,
                 self.file_id,
                 self.module_id,
                 krate,
@@ -1072,7 +1072,7 @@ pub fn collect_struct(
         }
     };
 
-    interner.set_doc_comments(ReferenceId::Struct(id), doc_comments);
+    interner.set_doc_comments(ReferenceId::Type(id), doc_comments);
 
     for (index, field) in unresolved.struct_def.fields.iter().enumerate() {
         if !field.doc_comments.is_empty() {
@@ -1106,7 +1106,7 @@ pub fn collect_struct(
     }
 
     if interner.is_in_lsp_mode() {
-        interner.register_struct(id, name.to_string(), visibility, parent_module_id);
+        interner.register_type(id, name.to_string(), visibility, parent_module_id);
     }
 
     Some((id, unresolved))
@@ -1167,7 +1167,7 @@ pub fn collect_enum(
         }
     };
 
-    interner.set_doc_comments(ReferenceId::Enum(id), doc_comments);
+    interner.set_doc_comments(ReferenceId::Type(id), doc_comments);
 
     for (index, variant) in unresolved.enum_def.variants.iter().enumerate() {
         if !variant.doc_comments.is_empty() {
@@ -1201,7 +1201,7 @@ pub fn collect_enum(
     }
 
     if interner.is_in_lsp_mode() {
-        interner.register_enum(id, name.to_string(), visibility, parent_module_id);
+        interner.register_type(id, name.to_string(), visibility, parent_module_id);
     }
 
     Some((id, unresolved))

compiler/noirc_frontend/src/hir/def_map/module_data.rs

@@ -31,7 +31,7 @@ pub struct ModuleData {
     /// True if this module is a `contract Foo { ... }` module containing contract functions
     pub is_contract: bool,
 
-    /// True if this module is actually a struct
+    /// True if this module is actually a type
     pub is_type: bool,
 
     pub attributes: Vec<SecondaryAttribute>,
@@ -44,7 +44,7 @@ impl ModuleData {
         outer_attributes: Vec<SecondaryAttribute>,
         inner_attributes: Vec<SecondaryAttribute>,
         is_contract: bool,
-        is_struct: bool,
+        is_type: bool,
     ) -> ModuleData {
         let mut attributes = outer_attributes;
         attributes.extend(inner_attributes);
@@ -57,7 +57,7 @@ impl ModuleData {
             definitions: ItemScope::default(),
             location,
             is_contract,
-            is_type: is_struct,
+            is_type,
             attributes,
         }
     }

compiler/noirc_frontend/src/hir_def/function.rs

@@ -141,6 +141,9 @@ pub struct FuncMeta {
     /// The trait impl this function belongs to, if any
     pub trait_impl: Option<TraitImplId>,
 
+    /// If this function is the one related to an enum variant, this holds its index (relative to `type_id`)
+    pub enum_variant_index: Option<usize>,
+
     /// True if this function is an entry point to the program.
     /// For non-contracts, this means the function is `main`.
     pub is_entry_point: bool,

compiler/noirc_frontend/src/hir_def/types.rs

@@ -448,13 +448,19 @@ impl DataType {
         self.body = TypeBody::Struct(fields);
     }
 
-    pub(crate) fn push_variant(&mut self, variant: EnumVariant) {
+    pub(crate) fn init_variants(&mut self) {
         match &mut self.body {
             TypeBody::None => {
-                self.body = TypeBody::Enum(vec![variant]);
+                self.body = TypeBody::Enum(vec![]);
             }
+            _ => panic!("Called init_variants but body was None"),
+        }
+    }
+
+    pub(crate) fn push_variant(&mut self, variant: EnumVariant) {
+        match &mut self.body {
             TypeBody::Enum(variants) => variants.push(variant),
-            TypeBody::Struct(_) => panic!("Called push_variant on a non-variant type {self}"),
+            _ => panic!("Called push_variant on {self} but body wasn't an enum"),
         }
     }
 

compiler/noirc_frontend/src/locations.rs

@@ -60,7 +60,7 @@ impl NodeInterner {
         match reference {
             ReferenceId::Module(id) => self.module_attributes(&id).location,
             ReferenceId::Function(id) => self.function_modifiers(&id).name_location,
-            ReferenceId::Struct(id) | ReferenceId::Enum(id) => {
+            ReferenceId::Type(id) => {
                 let typ = self.get_type(id);
                 let typ = typ.borrow();
                 Location::new(typ.name.span(), typ.location.file)
@@ -109,8 +109,8 @@ impl NodeInterner {
             ModuleDefId::FunctionId(func_id) => {
                 self.add_function_reference(func_id, location);
             }
-            ModuleDefId::TypeId(struct_id) => {
-                self.add_struct_reference(struct_id, location, is_self_type);
+            ModuleDefId::TypeId(type_id) => {
+                self.add_type_reference(type_id, location, is_self_type);
             }
             ModuleDefId::TraitId(trait_id) => {
                 self.add_trait_reference(trait_id, location, is_self_type);
@@ -128,13 +128,13 @@ impl NodeInterner {
         self.add_reference(ReferenceId::Module(id), location, false);
     }
 
-    pub(crate) fn add_struct_reference(
+    pub(crate) fn add_type_reference(
         &mut self,
         id: TypeId,
         location: Location,
         is_self_type: bool,
     ) {
-        self.add_reference(ReferenceId::Struct(id), location, is_self_type);
+        self.add_reference(ReferenceId::Type(id), location, is_self_type);
     }
 
     pub(crate) fn add_struct_member_reference(
@@ -326,25 +326,14 @@ impl NodeInterner {
         self.register_name_for_auto_import(name, ModuleDefId::GlobalId(id), visibility, None);
     }
 
-    pub(crate) fn register_struct(
+    pub(crate) fn register_type(
         &mut self,
         id: TypeId,
         name: String,
         visibility: ItemVisibility,
         parent_module_id: ModuleId,
     ) {
-        self.add_definition_location(ReferenceId::Struct(id), Some(parent_module_id));
-        self.register_name_for_auto_import(name, ModuleDefId::TypeId(id), visibility, None);
-    }
-
-    pub(crate) fn register_enum(
-        &mut self,
-        id: TypeId,
-        name: String,
-        visibility: ItemVisibility,
-        parent_module_id: ModuleId,
-    ) {
-        self.add_definition_location(ReferenceId::Enum(id), Some(parent_module_id));
+        self.add_definition_location(ReferenceId::Type(id), Some(parent_module_id));
         self.register_name_for_auto_import(name, ModuleDefId::TypeId(id), visibility, None);
     }
 

compiler/noirc_frontend/src/node_interner.rs

@@ -299,9 +299,8 @@ pub enum DependencyId {
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum ReferenceId {
     Module(ModuleId),
-    Struct(TypeId),
+    Type(TypeId),
     StructMember(TypeId, usize),
-    Enum(TypeId),
     EnumVariant(TypeId, usize),
     Trait(TraitId),
     Global(GlobalId),

tooling/lsp/src/modules.rs

@@ -16,7 +16,7 @@ pub(crate) fn module_def_id_to_reference_id(module_def_id: ModuleDefId) -> Refer
     match module_def_id {
         ModuleDefId::ModuleId(id) => ReferenceId::Module(id),
         ModuleDefId::FunctionId(id) => ReferenceId::Function(id),
-        ModuleDefId::TypeId(id) => ReferenceId::Struct(id),
+        ModuleDefId::TypeId(id) => ReferenceId::Type(id),
         ModuleDefId::TypeAliasId(id) => ReferenceId::Alias(id),
         ModuleDefId::TraitId(id) => ReferenceId::Trait(id),
         ModuleDefId::GlobalId(id) => ReferenceId::Global(id),

tooling/lsp/src/requests/code_action/fill_struct_fields.rs

@@ -20,25 +20,23 @@ impl<'a> CodeActionFinder<'a> {
         };
 
         let location = Location::new(path.span, self.file);
-        let Some(ReferenceId::Struct(struct_id)) = self.interner.find_referenced(location) else {
+        let Some(ReferenceId::Type(type_id)) = self.interner.find_referenced(location) else {
             return;
         };
 
-        let struct_type = self.interner.get_type(struct_id);
-        let struct_type = struct_type.borrow();
+        let typ = self.interner.get_type(type_id);
+        let typ = typ.borrow();
 
         // First get all of the struct's fields
-        let mut fields = struct_type.get_fields_as_written();
+        let Some(mut fields) = typ.get_fields_as_written() else {
+            return;
+        };
 
         // Remove the ones that already exists in the constructor
         for (constructor_field, _) in &constructor.fields {
             fields.retain(|field| field.name.0.contents != constructor_field.0.contents);
         }
 
-        if fields.is_empty() {
-            return;
-        }
-
         // Some fields are missing. Let's suggest a quick fix that adds them.
         let bytes = self.source.as_bytes();
         let right_brace_index = span.end() as usize - 1;

tooling/lsp/src/requests/completion.rs

@@ -199,16 +199,19 @@ impl<'a> NodeFinder<'a> {
         };
 
         let location = Location::new(span, self.file);
-        let Some(ReferenceId::Struct(struct_id)) = self.interner.find_referenced(location) else {
+        let Some(ReferenceId::Type(struct_id)) = self.interner.find_referenced(location) else {
             return;
         };
 
         let struct_type = self.interner.get_type(struct_id);
         let struct_type = struct_type.borrow();
 
         // First get all of the struct's fields
-        let mut fields: Vec<_> =
-            struct_type.get_fields_as_written().into_iter().enumerate().collect();
+        let Some(fields) = struct_type.get_fields_as_written() else {
+            return;
+        };
+
+        let mut fields = fields.into_iter().enumerate().collect::<Vec<_>>();
 
         // Remove the ones that already exists in the constructor
         for (used_name, _) in &constructor_expression.fields {
@@ -805,9 +808,11 @@ impl<'a> NodeFinder<'a> {
         prefix: &str,
         self_prefix: bool,
     ) {
-        for (field_index, (name, visibility, typ)) in
-            struct_type.get_fields_with_visibility(generics).iter().enumerate()
-        {
+        let Some(fields) = struct_type.get_fields_with_visibility(generics) else {
+            return;
+        };
+
+        for (field_index, (name, visibility, typ)) in fields.iter().enumerate() {
             if !struct_member_is_visible(struct_type.id, *visibility, self.module_id, self.def_maps)
             {
                 continue;
@@ -1953,8 +1958,7 @@ fn name_matches(name: &str, prefix: &str) -> bool {
 fn module_def_id_from_reference_id(reference_id: ReferenceId) -> Option<ModuleDefId> {
     match reference_id {
         ReferenceId::Module(module_id) => Some(ModuleDefId::ModuleId(module_id)),
-        ReferenceId::Struct(struct_id) => Some(ModuleDefId::TypeId(struct_id)),
-        ReferenceId::Enum(enum_id) => Some(ModuleDefId::TypeId(enum_id)),
+        ReferenceId::Type(struct_id) => Some(ModuleDefId::TypeId(struct_id)),
         ReferenceId::Trait(trait_id) => Some(ModuleDefId::TraitId(trait_id)),
         ReferenceId::Function(func_id) => Some(ModuleDefId::FunctionId(func_id)),
         ReferenceId::Alias(type_alias_id) => Some(ModuleDefId::TypeAliasId(type_alias_id)),