fix: Add comptime to trait impls and structs

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

@@ -173,7 +173,7 @@
 ///
 /// To:
 ///
 /// impl<Context> Storage<Contex> {
 ///    fn init(context: Context) -> Self {
 ///        Storage {
 ///             a_map: Map::new(context, 0, |context, slot| {
@@ -248,6 +248,7 @@
         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 @@
                 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,
@@ -509,7 +510,7 @@
                         }
                     }
                     TraitItem::Type { name } => {
                         // TODO(nickysn or alexvitkov): implement context.def_interner.push_empty_type_alias and get an id, instead of using TypeAliasId::dummy_id()
                         if let Err((first_def, second_def)) = self.def_collector.def_map.modules
                             [trait_id.0.local_id.0]
                             .declare_type_alias(name.clone(), TypeAliasId::dummy_id())
@@ -695,7 +696,7 @@
         //   if it's an inline module, or the first char of a the file if it's an external module.
         // - `location` will always point to the token "foo" in `mod foo` regardless of whether
         //   it's inline or external.
         // Eventually the location put in `ModuleData` is used for codelenses about `contract`s,
         // so we keep using `location` so that it continues to work as usual.
         let location = Location::new(mod_name.span(), mod_location.file);
         let new_module = ModuleData::new(parent, location, is_contract);

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

@@ -48,7 +48,7 @@
 
 use chumsky::prelude::*;
 use iter_extended::vecmap;
 use lalrpop_util::lalrpop_mod;
 use noirc_errors::{Span, Spanned};
 
 mod assertion;
@@ -62,8 +62,8 @@
 mod traits;
 mod types;
 
 // synthesized by LALRPOP
 lalrpop_mod!(pub noir_parser);
 
 #[cfg(test)]
 mod test_helpers;
@@ -87,7 +87,7 @@
 
     if cfg!(feature = "experimental_parser") {
         for parsed_item in &parsed_module.items {
             if lalrpop_parser_supports_kind(&parsed_item.kind) {
                 match &parsed_item.kind {
                     ItemKind::Import(parsed_use_tree) => {
                         prototype_parse_use_tree(Some(parsed_use_tree), source_program);
@@ -171,20 +171,8 @@
             .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 @@
 ///                    | 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 @@
 ///
 /// 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,
+            })
         })
 }