fix: Add comptime to trait impls and structs (#5511)

# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

## Summary\*

This is somewhat of a quick fix to the issue in the trait constraint PR
and in @asterite's comptime prefix operator PR.

The issue was that the comptime functions/blocks were using other
functions which weren't elaborated yet. So they saw an empty function
body. Just adding `comptime` to a function in an impl wasn't sufficient
since the compiler didn't look in every function within an impl to know
whether to add the impl to the list of comptime items. Moreover, it
wouldn't be possible to only add this `comptime` to a subset of
functions within an impl so I decided to allow `comptime` as a keyword
on an entire impl as well - and this is what decides whether or not to
add it to the comptime items.

I'm not really satisfied with this solution since it overloads
`comptime` to mean both "run this at compile-time (always)" and "_able
to be run_ at compile-time" but I don't have a better solution right
now. To some degree the separation is necessary since we need to
separate runtime functions that can be modified by comptime functions
from comptime functions that cannot be modified (because we've already
run them).

## Additional Context

Requires #5517 to work

## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [x] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

aztec_macros/src/transforms/note_interface.rs

@@ -73,6 +73,7 @@ pub fn generate_note_interface_impl(module: &mut SortedModule) -> Result<(), Azt
                 generics: vec![],
                 methods: vec![],
                 where_clause: vec![],
+                is_comptime: false,
             };
             module.impls.push(default_impl.clone());
             module.impls.last_mut().unwrap()

aztec_macros/src/transforms/storage.rs

@@ -248,6 +248,7 @@ pub fn generate_storage_implementation(
         methods: vec![(init, Span::default())],
 
         where_clause: vec![],
+        is_comptime: false,
     };
     module.impls.push(storage_impl);
 

compiler/noirc_frontend/src/ast/structure.rs

@@ -14,18 +14,7 @@ pub struct NoirStruct {
     pub generics: UnresolvedGenerics,
     pub fields: Vec<(Ident, UnresolvedType)>,
     pub span: Span,
-}
-
-impl NoirStruct {
-    pub fn new(
-        name: Ident,
-        attributes: Vec<SecondaryAttribute>,
-        generics: UnresolvedGenerics,
-        fields: Vec<(Ident, UnresolvedType)>,
-        span: Span,
-    ) -> NoirStruct {
-        NoirStruct { name, attributes, generics, fields, span }
-    }
+    pub is_comptime: bool,
 }
 
 impl Display for NoirStruct {

compiler/noirc_frontend/src/ast/traits.rs

@@ -53,6 +53,7 @@ pub struct TypeImpl {
     pub generics: UnresolvedGenerics,
     pub where_clause: Vec<UnresolvedTraitConstraint>,
     pub methods: Vec<(NoirFunction, Span)>,
+    pub is_comptime: bool,
 }
 
 /// Ast node for an implementation of a trait for a particular type
@@ -69,6 +70,8 @@ pub struct NoirTraitImpl {
     pub where_clause: Vec<UnresolvedTraitConstraint>,
 
     pub items: Vec<TraitImplItem>,
+
+    pub is_comptime: bool,
 }
 
 /// Represents a simple trait constraint such as `where Foo: TraitY<U, V>`

compiler/noirc_frontend/src/elaborator/mod.rs

@@ -1627,17 +1627,25 @@ impl<'context> Elaborator<'context> {
             function_sets.push(UnresolvedFunctions { functions, file_id, trait_id, self_type });
         }
 
+        let (comptime_trait_impls, trait_impls) =
+            items.trait_impls.into_iter().partition(|trait_impl| trait_impl.is_comptime);
+
+        let (comptime_structs, structs) =
+            items.types.into_iter().partition(|typ| typ.1.struct_def.is_comptime);
+
         let comptime = CollectedItems {
             functions: comptime_function_sets,
-            types: BTreeMap::new(),
+            types: comptime_structs,
             type_aliases: BTreeMap::new(),
             traits: BTreeMap::new(),
-            trait_impls: Vec::new(),
+            trait_impls: comptime_trait_impls,
             globals: Vec::new(),
             impls: rustc_hash::FxHashMap::default(),
         };
 
         items.functions = function_sets;
+        items.trait_impls = trait_impls;
+        items.types = structs;
         (comptime, items)
     }
 
@@ -1648,75 +1656,85 @@ impl<'context> Elaborator<'context> {
         location: Location,
     ) {
         for item in items {
-            match item {
-                TopLevelStatement::Function(function) => {
-                    let id = self.interner.push_empty_fn();
-                    let module = self.module_id();
-                    self.interner.push_function(id, &function.def, module, location);
-                    let functions = vec![(self.local_module, id, function)];
-                    generated_items.functions.push(UnresolvedFunctions {
-                        file_id: self.file,
-                        functions,
-                        trait_id: None,
-                        self_type: None,
-                    });
-                }
-                TopLevelStatement::TraitImpl(mut trait_impl) => {
-                    let methods = dc_mod::collect_trait_impl_functions(
-                        self.interner,
-                        &mut trait_impl,
-                        self.crate_id,
-                        self.file,
-                        self.local_module,
-                    );
+            self.add_item(item, generated_items, location);
+        }
+    }
 
-                    generated_items.trait_impls.push(UnresolvedTraitImpl {
-                        file_id: self.file,
-                        module_id: self.local_module,
-                        trait_generics: trait_impl.trait_generics,
-                        trait_path: trait_impl.trait_name,
-                        object_type: trait_impl.object_type,
-                        methods,
-                        generics: trait_impl.impl_generics,
-                        where_clause: trait_impl.where_clause,
-
-                        // These last fields are filled in later
-                        trait_id: None,
-                        impl_id: None,
-                        resolved_object_type: None,
-                        resolved_generics: Vec::new(),
-                        resolved_trait_generics: Vec::new(),
-                    });
-                }
-                TopLevelStatement::Global(global) => {
-                    let (global, error) = dc_mod::collect_global(
-                        self.interner,
-                        self.def_maps.get_mut(&self.crate_id).unwrap(),
-                        global,
-                        self.file,
-                        self.local_module,
-                    );
+    fn add_item(
+        &mut self,
+        item: TopLevelStatement,
+        generated_items: &mut CollectedItems,
+        location: Location,
+    ) {
+        match item {
+            TopLevelStatement::Function(function) => {
+                let id = self.interner.push_empty_fn();
+                let module = self.module_id();
+                self.interner.push_function(id, &function.def, module, location);
+                let functions = vec![(self.local_module, id, function)];
+                generated_items.functions.push(UnresolvedFunctions {
+                    file_id: self.file,
+                    functions,
+                    trait_id: None,
+                    self_type: None,
+                });
+            }
+            TopLevelStatement::TraitImpl(mut trait_impl) => {
+                let methods = dc_mod::collect_trait_impl_functions(
+                    self.interner,
+                    &mut trait_impl,
+                    self.crate_id,
+                    self.file,
+                    self.local_module,
+                );
 
-                    generated_items.globals.push(global);
-                    if let Some(error) = error {
-                        self.errors.push(error);
-                    }
-                }
-                // Assume that an error has already been issued
-                TopLevelStatement::Error => (),
-
-                TopLevelStatement::Module(_)
-                | TopLevelStatement::Import(_)
-                | TopLevelStatement::Struct(_)
-                | TopLevelStatement::Trait(_)
-                | TopLevelStatement::Impl(_)
-                | TopLevelStatement::TypeAlias(_)
-                | TopLevelStatement::SubModule(_) => {
-                    let item = item.to_string();
-                    let error = InterpreterError::UnsupportedTopLevelItemUnquote { item, location };
-                    self.errors.push(error.into_compilation_error_pair());
+                generated_items.trait_impls.push(UnresolvedTraitImpl {
+                    file_id: self.file,
+                    module_id: self.local_module,
+                    trait_generics: trait_impl.trait_generics,
+                    trait_path: trait_impl.trait_name,
+                    object_type: trait_impl.object_type,
+                    methods,
+                    generics: trait_impl.impl_generics,
+                    where_clause: trait_impl.where_clause,
+                    is_comptime: trait_impl.is_comptime,
+
+                    // These last fields are filled in later
+                    trait_id: None,
+                    impl_id: None,
+                    resolved_object_type: None,
+                    resolved_generics: Vec::new(),
+                    resolved_trait_generics: Vec::new(),
+                });
+            }
+            TopLevelStatement::Global(global) => {
+                let (global, error) = dc_mod::collect_global(
+                    self.interner,
+                    self.def_maps.get_mut(&self.crate_id).unwrap(),
+                    global,
+                    self.file,
+                    self.local_module,
+                );
+
+                generated_items.globals.push(global);
+                if let Some(error) = error {
+                    self.errors.push(error);
                 }
             }
+            // Assume that an error has already been issued
+            TopLevelStatement::Error => (),
+
+            TopLevelStatement::Module(_)
+            | TopLevelStatement::Import(_)
+            | TopLevelStatement::Struct(_)
+            | TopLevelStatement::Trait(_)
+            | TopLevelStatement::Impl(_)
+            | TopLevelStatement::TypeAlias(_)
+            | TopLevelStatement::SubModule(_) => {
+                let item = item.to_string();
+                let error = InterpreterError::UnsupportedTopLevelItemUnquote { item, location };
+                self.errors.push(error.into_compilation_error_pair());
+            }
         }
     }
 

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

@@ -124,6 +124,7 @@ pub struct UnresolvedTraitImpl {
     pub methods: UnresolvedFunctions,
     pub generics: UnresolvedGenerics,
     pub where_clause: Vec<UnresolvedTraitConstraint>,
+    pub is_comptime: bool,
 
     // Every field after this line is filled in later in the elaborator
     pub trait_id: Option<TraitId>,
@@ -279,7 +280,7 @@ impl DefCollector {
     /// Collect all of the definitions in a given crate into a CrateDefMap
     /// Modules which are not a part of the module hierarchy starting with
     /// the root module, will be ignored.
-    pub fn collect(
+    pub fn collect_crate_and_dependencies(
         mut def_map: CrateDefMap,
         context: &mut Context,
         ast: SortedModule,

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

@@ -185,6 +185,7 @@ impl<'a> ModCollector<'a> {
                 generics: trait_impl.impl_generics,
                 where_clause: trait_impl.where_clause,
                 trait_generics: trait_impl.trait_generics,
+                is_comptime: trait_impl.is_comptime,
 
                 // These last fields are filled later on
                 trait_id: None,

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

@@ -118,7 +118,7 @@ impl CrateDefMap {
         };
 
         // Now we want to populate the CrateDefMap using the DefCollector
-        errors.extend(DefCollector::collect(
+        errors.extend(DefCollector::collect_crate_and_dependencies(
             def_map,
             context,
             ast,

compiler/noirc_frontend/src/parser/mod.rs

@@ -39,6 +39,24 @@ pub enum TopLevelStatement {
     Error,
 }
 
+impl TopLevelStatement {
+    pub fn into_item_kind(self) -> Option<ItemKind> {
+        match self {
+            TopLevelStatement::Function(f) => Some(ItemKind::Function(f)),
+            TopLevelStatement::Module(m) => Some(ItemKind::ModuleDecl(m)),
+            TopLevelStatement::Import(i) => Some(ItemKind::Import(i)),
+            TopLevelStatement::Struct(s) => Some(ItemKind::Struct(s)),
+            TopLevelStatement::Trait(t) => Some(ItemKind::Trait(t)),
+            TopLevelStatement::TraitImpl(t) => Some(ItemKind::TraitImpl(t)),
+            TopLevelStatement::Impl(i) => Some(ItemKind::Impl(i)),
+            TopLevelStatement::TypeAlias(t) => Some(ItemKind::TypeAlias(t)),
+            TopLevelStatement::SubModule(s) => Some(ItemKind::Submodules(s)),
+            TopLevelStatement::Global(c) => Some(ItemKind::Global(c)),
+            TopLevelStatement::Error => None,
+        }
+    }
+}
+
 // Helper trait that gives us simpler type signatures for return types:
 // e.g. impl Parser<T> versus impl Parser<Token, T, Error = Simple<Token>>
 pub trait NoirParser<T>: Parser<Token, T, Error = ParserError> + Sized + Clone {}

compiler/noirc_frontend/src/parser/parser.rs

@@ -171,20 +171,8 @@ fn module() -> impl NoirParser<ParsedModule> {
             .to(ParsedModule::default())
             .then(spanned(top_level_statement(module_parser)).repeated())
             .foldl(|mut program, (statement, span)| {
-                let mut push_item = |kind| program.items.push(Item { kind, span });
-
-                match statement {
-                    TopLevelStatement::Function(f) => push_item(ItemKind::Function(f)),
-                    TopLevelStatement::Module(m) => push_item(ItemKind::ModuleDecl(m)),
-                    TopLevelStatement::Import(i) => push_item(ItemKind::Import(i)),
-                    TopLevelStatement::Struct(s) => push_item(ItemKind::Struct(s)),
-                    TopLevelStatement::Trait(t) => push_item(ItemKind::Trait(t)),
-                    TopLevelStatement::TraitImpl(t) => push_item(ItemKind::TraitImpl(t)),
-                    TopLevelStatement::Impl(i) => push_item(ItemKind::Impl(i)),
-                    TopLevelStatement::TypeAlias(t) => push_item(ItemKind::TypeAlias(t)),
-                    TopLevelStatement::SubModule(s) => push_item(ItemKind::Submodules(s)),
-                    TopLevelStatement::Global(c) => push_item(ItemKind::Global(c)),
-                    TopLevelStatement::Error => (),
+                if let Some(kind) = statement.into_item_kind() {
+                    program.items.push(Item { kind, span });
                 }
                 program
             })
@@ -204,9 +192,9 @@ pub fn top_level_items() -> impl NoirParser<Vec<TopLevelStatement>> {
 ///                    | module_declaration
 ///                    | use_statement
 ///                    | global_declaration
-fn top_level_statement(
-    module_parser: impl NoirParser<ParsedModule>,
-) -> impl NoirParser<TopLevelStatement> {
+fn top_level_statement<'a>(
+    module_parser: impl NoirParser<ParsedModule> + 'a,
+) -> impl NoirParser<TopLevelStatement> + 'a {
     choice((
         function::function_definition(false).map(TopLevelStatement::Function),
         structs::struct_definition(),
@@ -227,22 +215,24 @@ fn top_level_statement(
 ///
 /// implementation: 'impl' generics type '{' function_definition ... '}'
 fn implementation() -> impl NoirParser<TopLevelStatement> {
-    keyword(Keyword::Impl)
-        .ignore_then(function::generics())
+    maybe_comp_time()
+        .then_ignore(keyword(Keyword::Impl))
+        .then(function::generics())
         .then(parse_type().map_with_span(|typ, span| (typ, span)))
         .then(where_clause())
         .then_ignore(just(Token::LeftBrace))
         .then(spanned(function::function_definition(true)).repeated())
         .then_ignore(just(Token::RightBrace))
         .map(|args| {
             let ((other_args, where_clause), methods) = args;
-            let (generics, (object_type, type_span)) = other_args;
+            let ((is_comptime, generics), (object_type, type_span)) = other_args;
             TopLevelStatement::Impl(TypeImpl {
                 generics,
                 object_type,
                 type_span,
                 where_clause,
                 methods,
+                is_comptime,
             })
         })
 }

compiler/noirc_frontend/src/parser/parser/structs.rs

@@ -1,6 +1,7 @@
 use chumsky::prelude::*;
 
 use crate::ast::{Ident, NoirStruct, UnresolvedType};
+use crate::parser::parser::types::maybe_comp_time;
 use crate::{
     parser::{
         parser::{
@@ -28,13 +29,21 @@ pub(super) fn struct_definition() -> impl NoirParser<TopLevelStatement> {
         .or(just(Semicolon).to(Vec::new()));
 
     attributes()
+        .then(maybe_comp_time())
         .then_ignore(keyword(Struct))
         .then(ident())
         .then(function::generics())
         .then(fields)
-        .validate(|(((raw_attributes, name), generics), fields), span, emit| {
-            let attributes = validate_secondary_attributes(raw_attributes, span, emit);
-            TopLevelStatement::Struct(NoirStruct { name, attributes, generics, fields, span })
+        .validate(|((((attributes, is_comptime), name), generics), fields), span, emit| {
+            let attributes = validate_secondary_attributes(attributes, span, emit);
+            TopLevelStatement::Struct(NoirStruct {
+                name,
+                attributes,
+                generics,
+                fields,
+                span,
+                is_comptime,
+            })
         })
 }