feat(experimental): Support enums in comptime code (#7194)

noir-lang · Jan 27, 2025 · 4fe4cc6 · 4fe4cc6
1 parent c44b626
commit 4fe4cc6
Show file tree

Hide file tree

Showing 12 changed files with 88 additions and 14 deletions.
diff --git a/compiler/noirc_frontend/src/ast/expression.rs b/compiler/noirc_frontend/src/ast/expression.rs
@@ -739,8 +739,7 @@ impl Display for CastExpression {
 
 impl Display for ConstructorExpression {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let fields =
-            self.fields.iter().map(|(ident, expr)| format!("{ident}: {expr}")).collect::<Vec<_>>();
+        let fields = vecmap(&self.fields, |(ident, expr)| format!("{ident}: {expr}"));
 
         write!(f, "({} {{ {} }})", self.typ, fields.join(", "))
     }

diff --git a/compiler/noirc_frontend/src/elaborator/enums.rs b/compiler/noirc_frontend/src/elaborator/enums.rs
@@ -129,7 +129,6 @@ impl Elaborator<'_> {
         let enum_generics = self_type.borrow().generic_types();
         let construct_variant = HirExpression::EnumConstructor(HirEnumConstructorExpression {
             r#type: self_type.clone(),
-            enum_generics: enum_generics.clone(),
             arguments,
             variant_index,
         });

diff --git a/compiler/noirc_frontend/src/elaborator/mod.rs b/compiler/noirc_frontend/src/elaborator/mod.rs
@@ -1843,7 +1843,11 @@ impl<'context> Elaborator<'context> {
             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()));
+
+                // 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));
 
                 // Define a function for each variant to construct it
                 self.define_enum_variant_function(

diff --git a/compiler/noirc_frontend/src/hir/comptime/display.rs b/compiler/noirc_frontend/src/hir/comptime/display.rs
@@ -365,6 +365,22 @@ impl<'value, 'interner> Display for ValuePrinter<'value, 'interner> {
                 });
                 write!(f, "{typename} {{ {} }}", fields.join(", "))
             }
+            Value::Enum(tag, args, typ) => {
+                let args = vecmap(args, |arg| arg.display(self.interner).to_string()).join(", ");
+
+                match typ.follow_bindings_shallow().as_ref() {
+                    Type::DataType(def, _) => {
+                        let def = def.borrow();
+                        let variant = def.variant_at(*tag);
+                        if variant.is_function {
+                            write!(f, "{}::{}({args})", def.name, variant.name)
+                        } else {
+                            write!(f, "{}::{}", def.name, variant.name)
+                        }
+                    }
+                    other => write!(f, "{other}(args)"),
+                }
+            }
             Value::Pointer(value, _) => write!(f, "&mut {}", value.borrow().display(self.interner)),
             Value::Array(values, _) => {
                 let values = vecmap(values, |value| value.display(self.interner).to_string());

diff --git a/compiler/noirc_frontend/src/hir/comptime/errors.rs b/compiler/noirc_frontend/src/hir/comptime/errors.rs
@@ -133,6 +133,10 @@ pub enum InterpreterError {
         typ: Type,
         location: Location,
     },
+    NonEnumInConstructor {
+        typ: Type,
+        location: Location,
+    },
     CannotInlineMacro {
         value: String,
         typ: Type,
@@ -300,6 +304,7 @@ impl InterpreterError {
             | InterpreterError::CastToNonNumericType { location, .. }
             | InterpreterError::QuoteInRuntimeCode { location, .. }
             | InterpreterError::NonStructInConstructor { location, .. }
+            | InterpreterError::NonEnumInConstructor { location, .. }
             | InterpreterError::CannotInlineMacro { location, .. }
             | InterpreterError::UnquoteFoundDuringEvaluation { location, .. }
             | InterpreterError::UnsupportedTopLevelItemUnquote { location, .. }
@@ -505,6 +510,10 @@ impl<'a> From<&'a InterpreterError> for CustomDiagnostic {
                 let msg = format!("`{typ}` is not a struct type");
                 CustomDiagnostic::simple_error(msg, String::new(), location.span)
             }
+            InterpreterError::NonEnumInConstructor { typ, location } => {
+                let msg = format!("`{typ}` is not an enum type");
+                CustomDiagnostic::simple_error(msg, String::new(), location.span)
+            }
             InterpreterError::CannotInlineMacro { value, typ, location } => {
                 let msg = format!("Cannot inline values of type `{typ}` into this position");
                 let secondary = format!("Cannot inline value `{value}`");

diff --git a/compiler/noirc_frontend/src/hir/comptime/interpreter.rs b/compiler/noirc_frontend/src/hir/comptime/interpreter.rs
@@ -1290,10 +1290,12 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
 
     fn evaluate_enum_constructor(
         &mut self,
-        _constructor: HirEnumConstructorExpression,
-        _id: ExprId,
+        constructor: HirEnumConstructorExpression,
+        id: ExprId,
     ) -> IResult<Value> {
-        todo!("Support enums in the comptime interpreter")
+        let fields = try_vecmap(constructor.arguments, |arg| self.evaluate(arg))?;
+        let typ = self.elaborator.interner.id_type(id).follow_bindings();
+        Ok(Value::Enum(constructor.variant_index, fields, typ))
     }
 
     fn evaluate_access(&mut self, access: HirMemberAccess, id: ExprId) -> IResult<Value> {

diff --git a/compiler/noirc_frontend/src/hir/comptime/value.rs b/compiler/noirc_frontend/src/hir/comptime/value.rs
@@ -15,8 +15,8 @@ use crate::{
     elaborator::Elaborator,
     hir::{def_map::ModuleId, type_check::generics::TraitGenerics},
     hir_def::expr::{
-        HirArrayLiteral, HirConstructorExpression, HirExpression, HirIdent, HirLambda, HirLiteral,
-        ImplKind,
+        HirArrayLiteral, HirConstructorExpression, HirEnumConstructorExpression, HirExpression,
+        HirIdent, HirLambda, HirLiteral, ImplKind,
     },
     node_interner::{ExprId, FuncId, NodeInterner, StmtId, TraitId, TraitImplId, TypeId},
     parser::{Item, Parser},
@@ -55,6 +55,7 @@ pub enum Value {
 
     Tuple(Vec<Value>),
     Struct(HashMap<Rc<String>, Value>, Type),
+    Enum(/*tag*/ usize, /*args*/ Vec<Value>, Type),
     Pointer(Shared<Value>, /* auto_deref */ bool),
     Array(Vector<Value>, Type),
     Slice(Vector<Value>, Type),
@@ -131,6 +132,7 @@ impl Value {
                 Type::Tuple(vecmap(fields, |field| field.get_type().into_owned()))
             }
             Value::Struct(_, typ) => return Cow::Borrowed(typ),
+            Value::Enum(_, _, typ) => return Cow::Borrowed(typ),
             Value::Array(_, typ) => return Cow::Borrowed(typ),
             Value::Slice(_, typ) => return Cow::Borrowed(typ),
             Value::Quoted(_) => Type::Quoted(QuotedType::Quoted),
@@ -233,7 +235,7 @@ impl Value {
                     Ok((Ident::new(unwrap_rc(name), location.span), field))
                 })?;
 
-                let struct_type = match typ.follow_bindings() {
+                let struct_type = match typ.follow_bindings_shallow().as_ref() {
                     Type::DataType(def, _) => Some(def.borrow().id),
                     _ => return Err(InterpreterError::NonStructInConstructor { typ, location }),
                 };
@@ -246,6 +248,10 @@ impl Value {
                     struct_type,
                 }))
             }
+            value @ Value::Enum(..) => {
+                let hir = value.into_hir_expression(elaborator.interner, location)?;
+                ExpressionKind::Resolved(hir)
+            }
             Value::Array(elements, _) => {
                 let elements =
                     try_vecmap(elements, |element| element.into_expression(elaborator, location))?;
@@ -398,6 +404,21 @@ impl Value {
                     fields,
                 })
             }
+            Value::Enum(variant_index, args, typ) => {
+                let r#type = match typ.follow_bindings() {
+                    Type::DataType(def, _) => def,
+                    _ => return Err(InterpreterError::NonEnumInConstructor { typ, location }),
+                };
+
+                let arguments =
+                    try_vecmap(args, |arg| arg.into_hir_expression(interner, location))?;
+
+                HirExpression::EnumConstructor(HirEnumConstructorExpression {
+                    r#type,
+                    variant_index,
+                    arguments,
+                })
+            }
             Value::Array(elements, _) => {
                 let elements = try_vecmap(elements, |element| {
                     element.into_hir_expression(interner, location)

diff --git a/compiler/noirc_frontend/src/hir_def/expr.rs b/compiler/noirc_frontend/src/hir_def/expr.rs
@@ -297,7 +297,6 @@ pub struct HirConstructorExpression {
 #[derive(Debug, Clone)]
 pub struct HirEnumConstructorExpression {
     pub r#type: Shared<DataType>,
-    pub enum_generics: Vec<Type>,
     pub variant_index: usize,
 
     /// This refers to just the arguments that are passed. E.g. just

diff --git a/compiler/noirc_frontend/src/hir_def/types.rs b/compiler/noirc_frontend/src/hir_def/types.rs
@@ -368,11 +368,16 @@ pub struct StructField {
 pub struct EnumVariant {
     pub name: Ident,
     pub params: Vec<Type>,
+
+    /// True if this variant was declared as a function.
+    /// Required to distinguish `Foo::Bar` from `Foo::Bar()`
+    /// for zero-parameter variants. Only required for printing.
+    pub is_function: bool,
 }
 
 impl EnumVariant {
-    pub fn new(name: Ident, params: Vec<Type>) -> EnumVariant {
-        Self { name, params }
+    pub fn new(name: Ident, params: Vec<Type>, is_function: bool) -> EnumVariant {
+        Self { name, params, is_function }
     }
 }
 

diff --git a/test_programs/compile_success_empty/comptime_enums/Nargo.toml b/test_programs/compile_success_empty/comptime_enums/Nargo.toml
@@ -0,0 +1,6 @@
+[package]
+name = "comptime_enums"
+type = "bin"
+authors = [""]
+
+[dependencies]
diff --git a/test_programs/compile_success_empty/comptime_enums/src/main.nr b/test_programs/compile_success_empty/comptime_enums/src/main.nr
@@ -0,0 +1,13 @@
+fn main() {
+    comptime {
+        let _two = Foo::Couple(1, 2);
+        let _one = Foo::One(3);
+        let _none = Foo::None();
+    }
+}
+
+enum Foo {
+    Couple(i32, i32),
+    One(i32),
+    None,
+}
diff --git a/tooling/nargo_cli/build.rs b/tooling/nargo_cli/build.rs
@@ -66,13 +66,14 @@ const INLINER_MIN_OVERRIDES: [(&str, i64); 1] = [
 
 /// Some tests are expected to have warnings
 /// These should be fixed and removed from this list.
-const TESTS_WITH_EXPECTED_WARNINGS: [&str; 3] = [
+const TESTS_WITH_EXPECTED_WARNINGS: [&str; 4] = [
     // TODO(https://github.com/noir-lang/noir/issues/6238): remove from list once issue is closed
     "brillig_cast",
     // TODO(https://github.com/noir-lang/noir/issues/6238): remove from list once issue is closed
     "macros_in_comptime",
     // We issue a "experimental feature" warning for all enums until they're stabilized
     "enums",
+    "comptime_enums",
 ];
 
 fn read_test_cases(