feat: Parser and formatter support for enums (#7110)

compiler/noirc_frontend/src/ast/enumeration.rs

@@ -0,0 +1,50 @@
+use std::fmt::Display;
+
+use crate::ast::{Ident, UnresolvedGenerics, UnresolvedType};
+use crate::token::SecondaryAttribute;
+
+use iter_extended::vecmap;
+use noirc_errors::Span;
+
+use super::{Documented, ItemVisibility};
+
+/// Ast node for an enum
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct NoirEnumeration {
+    pub name: Ident,
+    pub attributes: Vec<SecondaryAttribute>,
+    pub visibility: ItemVisibility,
+    pub generics: UnresolvedGenerics,
+    pub variants: Vec<Documented<EnumVariant>>,
+    pub span: Span,
+}
+
+impl NoirEnumeration {
+    pub fn is_abi(&self) -> bool {
+        self.attributes.iter().any(|attr| attr.is_abi())
+    }
+}
+
+/// We only support  variants of the form `Name(A, B, ...)` currently.
+/// Enum variants like `Name { a: A, b: B, .. }` will be implemented later
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct EnumVariant {
+    pub name: Ident,
+    pub parameters: Vec<UnresolvedType>,
+}
+
+impl Display for NoirEnumeration {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let generics = vecmap(&self.generics, |generic| generic.to_string());
+        let generics = if generics.is_empty() { "".into() } else { generics.join(", ") };
+
+        writeln!(f, "enum {}{} {{", self.name, generics)?;
+
+        for variant in self.variants.iter() {
+            let parameters = vecmap(&variant.item.parameters, ToString::to_string).join(", ");
+            writeln!(f, "    {}({}),", variant.item.name, parameters)?;
+        }
+
+        write!(f, "}}")
+    }
+}

compiler/noirc_frontend/src/ast/mod.rs

@@ -5,6 +5,7 @@
 //! Noir's Ast is produced by the parser and taken as input to name resolution,
 //! where it is converted into the Hir (defined in the hir_def module).
 mod docs;
+mod enumeration;
 mod expression;
 mod function;
 mod statement;
@@ -24,6 +25,7 @@ use proptest_derive::Arbitrary;
 
 use acvm::FieldElement;
 pub use docs::*;
+pub use enumeration::*;
 use noirc_errors::Span;
 use serde::{Deserialize, Serialize};
 pub use statement::*;

compiler/noirc_frontend/src/ast/visitor.rs

@@ -21,15 +21,17 @@ use crate::{
 };
 
 use super::{
-    ForBounds, FunctionReturnType, GenericTypeArgs, IntegerBitSize, ItemVisibility, Pattern,
-    Signedness, TraitBound, TraitImplItemKind, TypePath, UnresolvedGenerics,
-    UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
+    ForBounds, FunctionReturnType, GenericTypeArgs, IntegerBitSize, ItemVisibility,
+    NoirEnumeration, Pattern, Signedness, TraitBound, TraitImplItemKind, TypePath,
+    UnresolvedGenerics, UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData,
+    UnresolvedTypeExpression,
 };
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub enum AttributeTarget {
     Module,
     Struct,
+    Enum,
     Trait,
     Function,
     Let,
@@ -142,6 +144,10 @@ pub trait Visitor {
         true
     }
 
+    fn visit_noir_enum(&mut self, _: &NoirEnumeration, _: Span) -> bool {
+        true
+    }
+
     fn visit_noir_type_alias(&mut self, _: &NoirTypeAlias, _: Span) -> bool {
         true
     }
@@ -527,6 +533,7 @@ impl Item {
             }
             ItemKind::TypeAlias(noir_type_alias) => noir_type_alias.accept(self.span, visitor),
             ItemKind::Struct(noir_struct) => noir_struct.accept(self.span, visitor),
+            ItemKind::Enum(noir_enum) => noir_enum.accept(self.span, visitor),
             ItemKind::ModuleDecl(module_declaration) => {
                 module_declaration.accept(self.span, visitor);
             }
@@ -775,6 +782,26 @@ impl NoirStruct {
     }
 }
 
+impl NoirEnumeration {
+    pub fn accept(&self, span: Span, visitor: &mut impl Visitor) {
+        if visitor.visit_noir_enum(self, span) {
+            self.accept_children(visitor);
+        }
+    }
+
+    pub fn accept_children(&self, visitor: &mut impl Visitor) {
+        for attribute in &self.attributes {
+            attribute.accept(AttributeTarget::Enum, visitor);
+        }
+
+        for variant in &self.variants {
+            for parameter in &variant.item.parameters {
+                parameter.accept(visitor);
+            }
+        }
+    }
+}
+
 impl NoirTypeAlias {
     pub fn accept(&self, span: Span, visitor: &mut impl Visitor) {
         if visitor.visit_noir_type_alias(self, span) {

compiler/noirc_frontend/src/elaborator/comptime.rs

@@ -442,7 +442,21 @@ impl<'context> Elaborator<'context> {
                     self.crate_id,
                     &mut self.errors,
                 ) {
-                    generated_items.types.insert(type_id, the_struct);
+                    generated_items.structs.insert(type_id, the_struct);
+                }
+            }
+            ItemKind::Enum(enum_def) => {
+                if let Some((type_id, the_enum)) = dc_mod::collect_enum(
+                    self.interner,
+                    self.def_maps.get_mut(&self.crate_id).unwrap(),
+                    self.usage_tracker,
+                    Documented::new(enum_def, item.doc_comments),
+                    self.file,
+                    self.local_module,
+                    self.crate_id,
+                    &mut self.errors,
+                ) {
+                    generated_items.enums.insert(type_id, the_enum);
                 }
             }
             ItemKind::Impl(r#impl) => {

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -319,7 +319,7 @@ impl<'context> Elaborator<'context> {
         }
 
         // Must resolve structs before we resolve globals.
-        self.collect_struct_definitions(&items.types);
+        self.collect_struct_definitions(&items.structs);
 
         self.define_function_metas(&mut items.functions, &mut items.impls, &mut items.trait_impls);
 
@@ -349,7 +349,7 @@ impl<'context> Elaborator<'context> {
         // since the generated items are checked beforehand as well.
         self.run_attributes(
             &items.traits,
-            &items.types,
+            &items.structs,
             &items.functions,
             &items.module_attributes,
         );

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

@@ -14,7 +14,7 @@ use crate::{Generics, Type};
 use crate::hir::resolution::import::{resolve_import, ImportDirective};
 use crate::hir::Context;
 
-use crate::ast::Expression;
+use crate::ast::{Expression, NoirEnumeration};
 use crate::node_interner::{
     FuncId, GlobalId, ModuleAttributes, NodeInterner, ReferenceId, StructId, TraitId, TraitImplId,
     TypeAliasId,
@@ -64,6 +64,12 @@ pub struct UnresolvedStruct {
     pub struct_def: NoirStruct,
 }
 
+pub struct UnresolvedEnum {
+    pub file_id: FileId,
+    pub module_id: LocalModuleId,
+    pub enum_def: NoirEnumeration,
+}
+
 #[derive(Clone)]
 pub struct UnresolvedTrait {
     pub file_id: FileId,
@@ -141,7 +147,8 @@ pub struct DefCollector {
 #[derive(Default)]
 pub struct CollectedItems {
     pub functions: Vec<UnresolvedFunctions>,
-    pub(crate) types: BTreeMap<StructId, UnresolvedStruct>,
+    pub(crate) structs: BTreeMap<StructId, UnresolvedStruct>,
+    pub(crate) enums: BTreeMap<StructId, UnresolvedEnum>,
     pub(crate) type_aliases: BTreeMap<TypeAliasId, UnresolvedTypeAlias>,
     pub(crate) traits: BTreeMap<TraitId, UnresolvedTrait>,
     pub globals: Vec<UnresolvedGlobal>,
@@ -153,7 +160,7 @@ pub struct CollectedItems {
 impl CollectedItems {
     pub fn is_empty(&self) -> bool {
         self.functions.is_empty()
-            && self.types.is_empty()
+            && self.structs.is_empty()
             && self.type_aliases.is_empty()
             && self.traits.is_empty()
             && self.globals.is_empty()
@@ -254,7 +261,8 @@ impl DefCollector {
             imports: vec![],
             items: CollectedItems {
                 functions: vec![],
-                types: BTreeMap::new(),
+                structs: BTreeMap::new(),
+                enums: BTreeMap::new(),
                 type_aliases: BTreeMap::new(),
                 traits: BTreeMap::new(),
                 impls: HashMap::default(),

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

@@ -12,8 +12,9 @@ use rustc_hash::FxHashMap as HashMap;
 
 use crate::ast::{
     Documented, Expression, FunctionDefinition, Ident, ItemVisibility, LetStatement,
-    ModuleDeclaration, NoirFunction, NoirStruct, NoirTrait, NoirTraitImpl, NoirTypeAlias, Pattern,
-    TraitImplItemKind, TraitItem, TypeImpl, UnresolvedType, UnresolvedTypeData,
+    ModuleDeclaration, NoirEnumeration, NoirFunction, NoirStruct, NoirTrait, NoirTraitImpl,
+    NoirTypeAlias, Pattern, TraitImplItemKind, TraitItem, TypeImpl, UnresolvedType,
+    UnresolvedTypeData,
 };
 use crate::hir::resolution::errors::ResolverError;
 use crate::node_interner::{ModuleAttributes, NodeInterner, ReferenceId, StructId};
@@ -27,8 +28,8 @@ use crate::{
 };
 use crate::{Generics, Kind, ResolvedGeneric, Type, TypeVariable};
 
-use super::dc_crate::CollectedItems;
 use super::dc_crate::ModuleAttribute;
+use super::dc_crate::{CollectedItems, UnresolvedEnum};
 use super::{
     dc_crate::{
         CompilationError, DefCollector, UnresolvedFunctions, UnresolvedGlobal, UnresolvedTraitImpl,
@@ -91,7 +92,7 @@ pub fn collect_defs(
 
     errors.extend(collector.collect_traits(context, ast.traits, crate_id));
 
-    errors.extend(collector.collect_structs(context, ast.types, crate_id));
+    errors.extend(collector.collect_structs(context, ast.structs, crate_id));
 
     errors.extend(collector.collect_type_aliases(context, ast.type_aliases, crate_id));
 
@@ -317,7 +318,7 @@ impl<'a> ModCollector<'a> {
                 krate,
                 &mut definition_errors,
             ) {
-                self.def_collector.items.types.insert(id, the_struct);
+                self.def_collector.items.structs.insert(id, the_struct);
             }
         }
         definition_errors
@@ -1078,6 +1079,20 @@ pub fn collect_struct(
     Some((id, unresolved))
 }
 
+#[allow(clippy::too_many_arguments)]
+pub fn collect_enum(
+    _interner: &mut NodeInterner,
+    _def_map: &mut CrateDefMap,
+    _usage_tracker: &mut UsageTracker,
+    _enum_definition: Documented<NoirEnumeration>,
+    _file_id: FileId,
+    _module_id: LocalModuleId,
+    _krate: CrateId,
+    _definition_errors: &mut [(CompilationError, FileId)],
+) -> Option<(StructId, UnresolvedEnum)> {
+    todo!("Implement collect_enum")
+}
+
 pub fn collect_impl(
     interner: &mut NodeInterner,
     items: &mut CollectedItems,

compiler/noirc_frontend/src/hir_def/types.rs

@@ -315,6 +315,7 @@ pub enum QuotedType {
     Type,
     TypedExpr,
     StructDefinition,
+    EnumDefinition,
     TraitConstraint,
     TraitDefinition,
     TraitImpl,
@@ -958,6 +959,7 @@ impl std::fmt::Display for QuotedType {
             QuotedType::Type => write!(f, "Type"),
             QuotedType::TypedExpr => write!(f, "TypedExpr"),
             QuotedType::StructDefinition => write!(f, "StructDefinition"),
+            QuotedType::EnumDefinition => write!(f, "EnumDefinition"),
             QuotedType::TraitDefinition => write!(f, "TraitDefinition"),
             QuotedType::TraitConstraint => write!(f, "TraitConstraint"),
             QuotedType::TraitImpl => write!(f, "TraitImpl"),

compiler/noirc_frontend/src/lexer/token.rs

@@ -1020,6 +1020,7 @@ pub enum Keyword {
     Dep,
     Else,
     Enum,
+    EnumDefinition,
     Expr,
     Field,
     Fn,
@@ -1080,6 +1081,7 @@ impl fmt::Display for Keyword {
             Keyword::Dep => write!(f, "dep"),
             Keyword::Else => write!(f, "else"),
             Keyword::Enum => write!(f, "enum"),
+            Keyword::EnumDefinition => write!(f, "EnumDefinition"),
             Keyword::Expr => write!(f, "Expr"),
             Keyword::Field => write!(f, "Field"),
             Keyword::Fn => write!(f, "fn"),
@@ -1143,6 +1145,7 @@ impl Keyword {
             "dep" => Keyword::Dep,
             "else" => Keyword::Else,
             "enum" => Keyword::Enum,
+            "EnumDefinition" => Keyword::EnumDefinition,
             "Expr" => Keyword::Expr,
             "Field" => Keyword::Field,
             "fn" => Keyword::Fn,

compiler/noirc_frontend/src/parser/errors.rs

@@ -83,8 +83,8 @@ pub enum ParserErrorReason {
         "Multiple primary attributes found. Only one function attribute is allowed per function"
     )]
     MultipleFunctionAttributesFound,
-    #[error("A function attribute cannot be placed on a struct")]
-    NoFunctionAttributesAllowedOnStruct,
+    #[error("A function attribute cannot be placed on a struct or enum")]
+    NoFunctionAttributesAllowedOnType,
     #[error("Assert statements can only accept string literals")]
     AssertMessageNotString,
     #[error("Integer bit size {0} isn't supported")]