feat(experimental): Implement enum tag constants (#7183)

Co-authored-by: Ary Borenszweig <asterite@gmail.com>

Author: jfecher

compiler/noirc_frontend/src/ast/enumeration.rs

@@ -30,7 +30,11 @@ impl NoirEnumeration {
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct EnumVariant {
     pub name: Ident,
-    pub parameters: Vec<UnresolvedType>,
+
+    /// This is None for tag variants without parameters.
+    /// A value of `Some(vec![])` corresponds to a variant defined as `Foo()`
+    /// with parenthesis but no parameters.
+    pub parameters: Option<Vec<UnresolvedType>>,
 }
 
 impl Display for NoirEnumeration {
@@ -41,8 +45,12 @@ impl Display for NoirEnumeration {
         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)?;
+            if let Some(parameters) = &variant.item.parameters {
+                let parameters = vecmap(parameters, ToString::to_string).join(", ");
+                writeln!(f, "    {}({}),", variant.item.name, parameters)?;
+            } else {
+                writeln!(f, "    {},", variant.item.name)?;
+            }
         }
 
         write!(f, "}}")

compiler/noirc_frontend/src/ast/visitor.rs

@@ -795,8 +795,10 @@ impl NoirEnumeration {
         }
 
         for variant in &self.variants {
-            for parameter in &variant.item.parameters {
-                parameter.accept(visitor);
+            if let Some(parameters) = &variant.item.parameters {
+                for parameter in parameters {
+                    parameter.accept(visitor);
+                }
             }
         }
     }

compiler/noirc_frontend/src/elaborator/enums.rs

@@ -8,16 +8,104 @@ use crate::{
         function::{FuncMeta, FunctionBody, HirFunction, Parameters},
         stmt::HirPattern,
     },
-    node_interner::{DefinitionKind, FuncId, FunctionModifiers, TypeId},
+    node_interner::{DefinitionKind, ExprId, FunctionModifiers, GlobalValue, TypeId},
     token::Attributes,
     DataType, Shared, Type,
 };
 
 use super::Elaborator;
 
 impl Elaborator<'_> {
+    /// Defines the value of an enum variant that we resolve an enum
+    /// variant expression to. E.g. `Foo::Bar` in `Foo::Bar(baz)`.
+    ///
+    /// If the variant requires arguments we should define a function,
+    /// otherwise we define a polymorphic global containing the tag value.
     #[allow(clippy::too_many_arguments)]
-    pub(super) fn define_enum_variant_function(
+    pub(super) fn define_enum_variant_constructor(
+        &mut self,
+        enum_: &NoirEnumeration,
+        type_id: TypeId,
+        variant: &EnumVariant,
+        variant_arg_types: Option<Vec<Type>>,
+        variant_index: usize,
+        datatype: &Shared<DataType>,
+        self_type: &Type,
+        self_type_unresolved: UnresolvedType,
+    ) {
+        match variant_arg_types {
+            Some(args) => self.define_enum_variant_function(
+                enum_,
+                type_id,
+                variant,
+                args,
+                variant_index,
+                datatype,
+                self_type,
+                self_type_unresolved,
+            ),
+            None => self.define_enum_variant_global(
+                enum_,
+                type_id,
+                variant,
+                variant_index,
+                datatype,
+                self_type,
+            ),
+        }
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn define_enum_variant_global(
+        &mut self,
+        enum_: &NoirEnumeration,
+        type_id: TypeId,
+        variant: &EnumVariant,
+        variant_index: usize,
+        datatype: &Shared<DataType>,
+        self_type: &Type,
+    ) {
+        let name = &variant.name;
+        let location = Location::new(variant.name.span(), self.file);
+
+        let global_id = self.interner.push_empty_global(
+            name.clone(),
+            type_id.local_module_id(),
+            type_id.krate(),
+            self.file,
+            Vec::new(),
+            false,
+            false,
+        );
+
+        let mut typ = self_type.clone();
+        if !datatype.borrow().generics.is_empty() {
+            let typevars = vecmap(&datatype.borrow().generics, |generic| generic.type_var.clone());
+            typ = Type::Forall(typevars, Box::new(typ));
+        }
+
+        let definition_id = self.interner.get_global(global_id).definition_id;
+        self.interner.push_definition_type(definition_id, typ.clone());
+
+        let no_parameters = Parameters(Vec::new());
+        let global_body =
+            self.make_enum_variant_constructor(datatype, variant_index, &no_parameters, location);
+        let let_statement = crate::hir_def::stmt::HirStatement::Expression(global_body);
+
+        let statement_id = self.interner.get_global(global_id).let_statement;
+        self.interner.replace_statement(statement_id, let_statement);
+
+        self.interner.get_global_mut(global_id).value = GlobalValue::Resolved(
+            crate::hir::comptime::Value::Enum(variant_index, Vec::new(), typ),
+        );
+
+        Self::get_module_mut(self.def_maps, type_id.module_id())
+            .declare_global(name.clone(), enum_.visibility, global_id)
+            .ok();
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    fn define_enum_variant_function(
         &mut self,
         enum_: &NoirEnumeration,
         type_id: TypeId,
@@ -48,7 +136,10 @@ impl Elaborator<'_> {
 
         let hir_name = HirIdent::non_trait_method(definition_id, location);
         let parameters = self.make_enum_variant_parameters(variant_arg_types, location);
-        self.push_enum_variant_function_body(id, datatype, variant_index, &parameters, location);
+
+        let body =
+            self.make_enum_variant_constructor(datatype, variant_index, &parameters, location);
+        self.interner.update_fn(id, HirFunction::unchecked_from_expr(body));
 
         let function_type =
             datatype_ref.variant_function_type_with_forall(variant_index, datatype.clone());
@@ -106,14 +197,13 @@ impl Elaborator<'_> {
     //     }
     // }
     // ```
-    fn push_enum_variant_function_body(
+    fn make_enum_variant_constructor(
         &mut self,
-        id: FuncId,
         self_type: &Shared<DataType>,
         variant_index: usize,
         parameters: &Parameters,
         location: Location,
-    ) {
+    ) -> ExprId {
         // Each parameter of the enum variant function is used as a parameter of the enum
         // constructor expression
         let arguments = vecmap(&parameters.0, |(pattern, typ, _)| match pattern {
@@ -126,18 +216,18 @@ impl Elaborator<'_> {
             _ => unreachable!(),
         });
 
-        let enum_generics = self_type.borrow().generic_types();
-        let construct_variant = HirExpression::EnumConstructor(HirEnumConstructorExpression {
+        let constructor = HirExpression::EnumConstructor(HirEnumConstructorExpression {
             r#type: self_type.clone(),
             arguments,
             variant_index,
         });
-        let body = self.interner.push_expr(construct_variant);
-        self.interner.update_fn(id, HirFunction::unchecked_from_expr(body));
 
+        let body = self.interner.push_expr(constructor);
+        let enum_generics = self_type.borrow().generic_types();
         let typ = Type::DataType(self_type.clone(), enum_generics);
         self.interner.push_expr_type(body, typ);
         self.interner.push_expr_location(body, location.span, location.file);
+        body
     }
 
     fn make_enum_variant_parameters(

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -1841,16 +1841,16 @@ impl<'context> Elaborator<'context> {
             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 parameters = variant.item.parameters.as_ref();
+                let types =
+                    parameters.map(|params| vecmap(params, |typ| self.resolve_type(typ.clone())));
                 let name = variant.item.name.clone();
 
-                // false here is for the eventual change to allow enum "constants" rather than
-                // always having them be called as functions. This can be replaced with an actual
-                // check once #7172 is implemented.
-                datatype.borrow_mut().push_variant(EnumVariant::new(name, types.clone(), false));
+                let is_function = types.is_some();
+                let params = types.clone().unwrap_or_default();
+                datatype.borrow_mut().push_variant(EnumVariant::new(name, params, is_function));
 
-                // Define a function for each variant to construct it
-                self.define_enum_variant_function(
+                self.define_enum_variant_constructor(
                     &typ.enum_def,
                     *type_id,
                     &variant.item,

compiler/noirc_frontend/src/hir/comptime/interpreter.rs

@@ -1294,7 +1294,7 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
         id: ExprId,
     ) -> IResult<Value> {
         let fields = try_vecmap(constructor.arguments, |arg| self.evaluate(arg))?;
-        let typ = self.elaborator.interner.id_type(id).follow_bindings();
+        let typ = self.elaborator.interner.id_type(id).unwrap_forall().1.follow_bindings();
         Ok(Value::Enum(constructor.variant_index, fields, typ))
     }
 

compiler/noirc_frontend/src/hir/comptime/value.rs

@@ -405,7 +405,8 @@ impl Value {
                 })
             }
             Value::Enum(variant_index, args, typ) => {
-                let r#type = match typ.follow_bindings() {
+                // Enum constants can have generic types but aren't functions
+                let r#type = match typ.unwrap_forall().1.follow_bindings() {
                     Type::DataType(def, _) => def,
                     _ => return Err(InterpreterError::NonEnumInConstructor { typ, location }),
                 };

compiler/noirc_frontend/src/monomorphization/mod.rs

@@ -2209,7 +2209,7 @@ fn unwrap_enum_type(
     typ: &HirType,
     location: Location,
 ) -> Result<Vec<(String, Vec<HirType>)>, MonomorphizationError> {
-    match typ.follow_bindings() {
+    match typ.unwrap_forall().1.follow_bindings() {
         HirType::DataType(def, args) => {
             // Some of args might not be mentioned in fields, so we need to check that they aren't unbound.
             for arg in &args {

compiler/noirc_frontend/src/parser/parser/enums.rs

@@ -92,12 +92,10 @@ impl<'a> Parser<'a> {
             self.bump();
         }
 
-        let mut parameters = Vec::new();
-
-        if self.eat_left_paren() {
+        let parameters = self.eat_left_paren().then(|| {
             let comma_separated = separated_by_comma_until_right_paren();
-            parameters = self.parse_many("variant parameters", comma_separated, Self::parse_type);
-        }
+            self.parse_many("variant parameters", comma_separated, Self::parse_type)
+        });
 
         Some(Documented::new(EnumVariant { name, parameters }, doc_comments))
     }
@@ -189,18 +187,19 @@ mod tests {
         let variant = noir_enum.variants.remove(0).item;
         assert_eq!("X", variant.name.to_string());
         assert!(matches!(
-            variant.parameters[0].typ,
+            variant.parameters.as_ref().unwrap()[0].typ,
             UnresolvedTypeData::Integer(Signedness::Signed, IntegerBitSize::ThirtyTwo)
         ));
 
         let variant = noir_enum.variants.remove(0).item;
         assert_eq!("y", variant.name.to_string());
-        assert!(matches!(variant.parameters[0].typ, UnresolvedTypeData::FieldElement));
-        assert!(matches!(variant.parameters[1].typ, UnresolvedTypeData::Integer(..)));
+        let parameters = variant.parameters.as_ref().unwrap();
+        assert!(matches!(parameters[0].typ, UnresolvedTypeData::FieldElement));
+        assert!(matches!(parameters[1].typ, UnresolvedTypeData::Integer(..)));
 
         let variant = noir_enum.variants.remove(0).item;
         assert_eq!("Z", variant.name.to_string());
-        assert_eq!(variant.parameters.len(), 0);
+        assert!(variant.parameters.is_none());
     }
 
     #[test]

test_programs/compile_success_empty/comptime_enums/src/main.nr

@@ -2,7 +2,7 @@ fn main() {
     comptime {
         let _two = Foo::Couple(1, 2);
         let _one = Foo::One(3);
-        let _none = Foo::None();
+        let _none = Foo::None;
     }
 }
 

test_programs/compile_success_empty/enums/src/main.nr

@@ -3,9 +3,10 @@ fn main() {
     let _b: Foo<u16> = Foo::B(3);
     let _c = Foo::C(4);
 
-    // (#7172): Single variant enums must be called as functions currently
     let _d: fn() -> Foo<(i32, i32)> = Foo::D;
     let _d: Foo<(i32, i32)> = Foo::D();
+    let _e: Foo<u16> = Foo::E;
+    let _e: Foo<u32> = Foo::E; // Ensure we can still use Foo::E polymorphically
 
     // Enum variants are functions and can be passed around as such
     let _many_cs = [1, 2, 3].map(Foo::C);
@@ -15,5 +16,6 @@ enum Foo<T> {
     A(Field, Field),
     B(u32),
     C(T),
-    D,
+    D(),
+    E,
 }