chore(experimental): Prevent enum panics by returning Options where p…

…ossible instead of panicking (#7180)

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

compiler/noirc_driver/src/abi_gen.rs

@@ -112,7 +112,7 @@ pub(super) fn abi_type_from_hir_type(context: &Context, typ: &Type) -> AbiType {
 
         Type::DataType(def, args) => {
             let struct_type = def.borrow();
-            let fields = struct_type.get_fields(args);
+            let fields = struct_type.get_fields(args).unwrap_or_default();
             let fields =
                 vecmap(fields, |(name, typ)| (name, abi_type_from_hir_type(context, &typ)));
             // For the ABI, we always want to resolve the struct paths from the root crate

compiler/noirc_driver/src/lib.rs

@@ -551,9 +551,10 @@ fn compile_contract_inner(
                     .map(|struct_id| {
                         let typ = context.def_interner.get_type(struct_id);
                         let typ = typ.borrow();
-                        let fields = vecmap(typ.get_fields(&[]), |(name, typ)| {
-                            (name, abi_type_from_hir_type(context, &typ))
-                        });
+                        let fields =
+                            vecmap(typ.get_fields(&[]).unwrap_or_default(), |(name, typ)| {
+                                (name, abi_type_from_hir_type(context, &typ))
+                            });
                         let path =
                             context.fully_qualified_struct_path(context.root_crate_id(), typ.id);
                         AbiType::Struct { path, fields }

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

@@ -622,7 +622,9 @@ impl<'context> Elaborator<'context> {
             (typ, generics)
         } else {
             match self.lookup_type_or_error(path) {
-                Some(Type::DataType(r#type, struct_generics)) => (r#type, struct_generics),
+                Some(Type::DataType(r#type, struct_generics)) if r#type.borrow().is_struct() => {
+                    (r#type, struct_generics)
+                }
                 Some(typ) => {
                     self.push_err(ResolverError::NonStructUsedInConstructor {
                         typ: typ.to_string(),
@@ -649,7 +651,11 @@ impl<'context> Elaborator<'context> {
         let generics = struct_generics.clone();
 
         let fields = constructor.fields;
-        let field_types = r#type.borrow().get_fields_with_visibility(&struct_generics);
+        let field_types = r#type
+            .borrow()
+            .get_fields_with_visibility(&struct_generics)
+            .expect("This type should already be validated to be a struct");
+
         let fields =
             self.resolve_constructor_expr_fields(struct_type.clone(), field_types, fields, span);
         let expr = HirExpression::Constructor(HirConstructorExpression {
@@ -660,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))
     }
@@ -683,7 +689,10 @@ impl<'context> Elaborator<'context> {
     ) -> Vec<(Ident, ExprId)> {
         let mut ret = Vec::with_capacity(fields.len());
         let mut seen_fields = HashSet::default();
-        let mut unseen_fields = struct_type.borrow().field_names();
+        let mut unseen_fields = struct_type
+            .borrow()
+            .field_names()
+            .expect("This type should already be validated to be a struct");
 
         for (field_name, field) in fields {
             let expected_field_with_index = field_types

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -12,25 +12,22 @@ 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},
-            errors::DefCollectorErrorKind,
+            filter_literal_globals, CollectedItems, CompilationError, ImplMap, UnresolvedEnum,
+            UnresolvedFunctions, UnresolvedGlobal, UnresolvedStruct, UnresolvedTraitImpl,
+            UnresolvedTypeAlias,
         },
+        def_collector::errors::DefCollectorErrorKind,
         def_map::{DefMaps, ModuleData},
         def_map::{LocalModuleId, 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 +1003,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(),
@@ -1035,13 +1033,21 @@ impl<'context> Elaborator<'context> {
                     self.mark_type_as_used(typ);
                 }
             }
-            Type::DataType(struct_type, generics) => {
-                self.mark_struct_as_constructed(struct_type.clone());
+            Type::DataType(datatype, generics) => {
+                self.mark_struct_as_constructed(datatype.clone());
                 for generic in generics {
                     self.mark_type_as_used(generic);
                 }
-                for (_, typ) in struct_type.borrow().get_fields(generics) {
-                    self.mark_type_as_used(&typ);
+                if let Some(fields) = datatype.borrow().get_fields(generics) {
+                    for (_, typ) in fields {
+                        self.mark_type_as_used(&typ);
+                    }
+                } else if let Some(variants) = datatype.borrow().get_variants(generics) {
+                    for (_, variant_types) in variants {
+                        for typ in variant_types {
+                            self.mark_type_as_used(&typ);
+                        }
+                    }
                 }
             }
             Type::Alias(alias_type, generics) => {
@@ -1776,7 +1782,7 @@ impl<'context> Elaborator<'context> {
             // Only handle structs without generics as any generics args will be checked
             // after monomorphization when performing SSA codegen
             if struct_type.borrow().generics.is_empty() {
-                let fields = struct_type.borrow().get_fields(&[]);
+                let fields = struct_type.borrow().get_fields(&[]).unwrap();
                 for (_, field_type) in fields.iter() {
                     if field_type.is_nested_slice() {
                         let location = struct_type.borrow().location;
@@ -1831,6 +1837,9 @@ impl<'context> Elaborator<'context> {
                 span: typ.enum_def.span,
             };
 
+            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();
@@ -1848,7 +1857,8 @@ impl<'context> Elaborator<'context> {
                     unresolved.clone(),
                 );
 
-                self.interner.add_definition_location(ReferenceId::EnumVariant(*type_id, i), None);
+                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

@@ -192,7 +192,11 @@ impl<'context> Elaborator<'context> {
         };
 
         let (struct_type, generics) = match self.lookup_type_or_error(name) {
-            Some(Type::DataType(struct_type, struct_generics)) => (struct_type, struct_generics),
+            Some(Type::DataType(struct_type, struct_generics))
+                if struct_type.borrow().is_struct() =>
+            {
+                (struct_type, struct_generics)
+            }
             None => return error_identifier(self),
             Some(typ) => {
                 let typ = typ.to_string();
@@ -234,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);
@@ -260,7 +264,10 @@ impl<'context> Elaborator<'context> {
     ) -> Vec<(Ident, HirPattern)> {
         let mut ret = Vec::with_capacity(fields.len());
         let mut seen_fields = HashSet::default();
-        let mut unseen_fields = struct_type.borrow().field_names();
+        let mut unseen_fields = struct_type
+            .borrow()
+            .field_names()
+            .expect("This type should already be validated to be a struct");
 
         for (field, pattern) in fields {
             let (field_type, visibility) = expected_type

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,
@@ -1477,7 +1477,7 @@ impl<'context> Elaborator<'context> {
             let datatype = datatype.borrow();
             let mut has_field_with_function_type = false;
 
-            if let Some(fields) = datatype.try_fields_raw() {
+            if let Some(fields) = datatype.fields_raw() {
                 has_field_with_function_type = fields
                     .iter()
                     .any(|field| field.name.0.contents == method_name && field.typ.is_function());

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

@@ -5,8 +5,8 @@ use crate::ast::{
     ArrayLiteral, AssignStatement, BlockExpression, CallExpression, CastExpression, ConstrainKind,
     ConstructorExpression, ExpressionKind, ForLoopStatement, ForRange, GenericTypeArgs, Ident,
     IfExpression, IndexExpression, InfixExpression, LValue, Lambda, Literal,
-    MemberAccessExpression, MethodCallExpression, Path, PathSegment, Pattern, PrefixExpression,
-    UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
+    MemberAccessExpression, MethodCallExpression, Path, PathKind, PathSegment, Pattern,
+    PrefixExpression, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
 };
 use crate::ast::{ConstrainStatement, Expression, Statement, StatementKind};
 use crate::hir_def::expr::{
@@ -217,7 +217,9 @@ impl HirExpression {
             HirExpression::EnumConstructor(constructor) => {
                 let typ = constructor.r#type.borrow();
                 let variant = &typ.variant_at(constructor.variant_index);
-                let path = Path::from_single(variant.name.to_string(), span);
+                let segment1 = PathSegment { ident: typ.name.clone(), span, generics: None };
+                let segment2 = PathSegment { ident: variant.name.clone(), span, generics: None };
+                let path = Path { segments: vec![segment1, segment2], kind: PathKind::Plain, span };
                 let func = Box::new(Expression::new(ExpressionKind::Variable(path), span));
                 let arguments = vecmap(&constructor.arguments, |arg| arg.to_display_ast(interner));
                 let call = CallExpression { func, arguments, is_macro_call: false };

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

@@ -569,9 +569,11 @@ fn struct_def_fields(
 
     let mut fields = im::Vector::new();
 
-    for (field_name, field_type) in struct_def.get_fields(&generic_args) {
-        let name = Value::Quoted(Rc::new(vec![Token::Ident(field_name)]));
-        fields.push_back(Value::Tuple(vec![name, Value::Type(field_type)]));
+    if let Some(struct_fields) = struct_def.get_fields(&generic_args) {
+        for (field_name, field_type) in struct_fields {
+            let name = Value::Quoted(Rc::new(vec![Token::Ident(field_name)]));
+            fields.push_back(Value::Tuple(vec![name, Value::Type(field_type)]));
+        }
     }
 
     let typ = Type::Slice(Box::new(Type::Tuple(vec![
@@ -597,10 +599,12 @@ fn struct_def_fields_as_written(
 
     let mut fields = im::Vector::new();
 
-    for field in struct_def.get_fields_as_written() {
-        let name = Value::Quoted(Rc::new(vec![Token::Ident(field.name.to_string())]));
-        let typ = Value::Type(field.typ);
-        fields.push_back(Value::Tuple(vec![name, typ]));
+    if let Some(struct_fields) = struct_def.get_fields_as_written() {
+        for field in struct_fields {
+            let name = Value::Quoted(Rc::new(vec![Token::Ident(field.name.to_string())]));
+            let typ = Value::Type(field.typ);
+            fields.push_back(Value::Tuple(vec![name, typ]));
+        }
     }
 
     let typ = Type::Slice(Box::new(Type::Tuple(vec![
@@ -1456,7 +1460,8 @@ fn zeroed(return_type: Type, span: Span) -> IResult<Value> {
         Type::Unit => Ok(Value::Unit),
         Type::Tuple(fields) => Ok(Value::Tuple(try_vecmap(fields, |field| zeroed(field, span))?)),
         Type::DataType(struct_type, generics) => {
-            let fields = struct_type.borrow().get_fields(&generics);
+            // TODO: Handle enums
+            let fields = struct_type.borrow().get_fields(&generics).unwrap();
             let mut values = HashMap::default();
 
             for (field_name, field_type) in fields {

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

@@ -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/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,