feat: Macros for initializer checks

noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr

@@ -18,15 +18,15 @@ contract StatefulTest {
     }
 
     #[aztec(private)]
+    #[aztec(initializer)]
     fn constructor(owner: AztecAddress, value: Field) {
         let selector = FunctionSelector::from_signature("create_note_no_init_check((Field),Field)");
         let _res = context.call_private_function(context.this_address(), selector, [owner.to_field(), value]);
-        mark_as_initialized(&mut context);
     }
 
     #[aztec(private)]
+    #[aztec(initcheck)]
     fn create_note(owner: AztecAddress, value: Field) {
-        assert_is_initialized(&mut context);
         if (value != 0) {
             let loc = storage.notes.at(owner);
             increment(loc, value, owner);

noir/noir-repo/aztec_macros/src/lib.rs

@@ -74,6 +74,7 @@ pub enum AztecMacroError {
     UnsupportedStorageType { span: Option<Span>, typ: UnresolvedTypeData },
     CouldNotAssignStorageSlots { secondary_message: Option<String> },
     EventError { span: Span, message: String },
+    UnsupportedAttributes { span: Span, secondary_message: Option<String> },
 }
 
 impl From<AztecMacroError> for MacroError {
@@ -114,6 +115,11 @@ impl From<AztecMacroError> for MacroError {
                 secondary_message: None,
                 span: Some(span),
             },
+            AztecMacroError::UnsupportedAttributes { span, secondary_message } => MacroError {
+                primary_message: "Unsupported attributes in contract function".to_string(),
+                secondary_message,
+                span: Some(span),
+            },
         }
     }
 }
@@ -430,22 +436,43 @@ fn transform_module(
     }
 
     for func in module.functions.iter_mut() {
+        let mut is_private = false;
+        let mut is_public = false;
+        let mut is_public_vm = false;
+        let mut is_initializer = false;
+        let mut skip_init_check = true; // Default to true once we're confident that the approach works
+
         for secondary_attribute in func.def.attributes.secondary.clone() {
-            let crate_graph = &context.crate_graph[crate_id];
             if is_custom_attribute(&secondary_attribute, "aztec(private)") {
-                transform_function("Private", func, storage_defined)
-                    .map_err(|err| (err, crate_graph.root_file_id))?;
-                has_transformed_module = true;
+                is_private = true;
+            } else if is_custom_attribute(&secondary_attribute, "aztec(initializer)") {
+                is_initializer = true;
+            } else if is_custom_attribute(&secondary_attribute, "aztec(initcheck)") {
+                skip_init_check = false;
             } else if is_custom_attribute(&secondary_attribute, "aztec(public)") {
-                transform_function("Public", func, storage_defined)
-                    .map_err(|err| (err, crate_graph.root_file_id))?;
-                has_transformed_module = true;
+                is_public = true;
             } else if is_custom_attribute(&secondary_attribute, "aztec(public-vm)") {
-                transform_vm_function(func, storage_defined)
-                    .map_err(|err| (err, crate_graph.root_file_id))?;
-                has_transformed_module = true;
+                is_public_vm = true;
             }
         }
+
+        // Apply transformations to the function based on collected attributes
+        if is_private || is_public {
+            transform_function(
+                if is_private { "Private" } else { "Public" },
+                func,
+                storage_defined,
+                is_initializer,
+                skip_init_check,
+            )
+            .map_err(|err| (err, crate_graph.root_file_id))?;
+            has_transformed_module = true;
+        } else if is_public_vm {
+            transform_vm_function(func, storage_defined)
+                .map_err(|err| (err, crate_graph.root_file_id))?;
+            has_transformed_module = true;
+        }
+
         // Add the storage struct to the beginning of the function if it is unconstrained in an aztec contract
         if storage_defined && func.def.is_unconstrained {
             transform_unconstrained(func);
@@ -627,11 +654,28 @@ fn transform_function(
     ty: &str,
     func: &mut NoirFunction,
     storage_defined: bool,
+    is_initializer: bool,
+    skip_init_check: bool,
 ) -> Result<(), AztecMacroError> {
     let context_name = format!("{}Context", ty);
     let inputs_name = format!("{}ContextInputs", ty);
     let return_type_name = format!("{}CircuitPublicInputs", ty);
 
+    // Add initialization check
+    if !skip_init_check {
+        if ty == "Public" {
+            let error = AztecMacroError::UnsupportedAttributes {
+                span: func.def.name.span(),
+                secondary_message: Some(
+                    "public functions do not yet support initialization check".to_owned(),
+                ),
+            };
+            return Err(error);
+        }
+        let init_check = create_init_check();
+        func.def.body.0.insert(0, init_check);
+    }
+
     // Add access to the storage struct
     if storage_defined {
         let storage_def = abstract_storage(&ty.to_lowercase(), false);
@@ -655,6 +699,21 @@ fn transform_function(
         func.def.body.0.push(return_values);
     }
 
+    // Before returning mark the contract as initialized
+    if is_initializer {
+        if ty == "Public" {
+            let error = AztecMacroError::UnsupportedAttributes {
+                span: func.def.name.span(),
+                secondary_message: Some(
+                    "public functions cannot yet be used as initializers".to_owned(),
+                ),
+            };
+            return Err(error);
+        }
+        let mark_initialized = create_mark_as_initialized();
+        func.def.body.0.push(mark_initialized);
+    }
+
     // Push the finish method call to the end of the function
     let finish_def = create_context_finish();
     func.def.body.0.push(finish_def);
@@ -1091,6 +1150,32 @@ fn create_inputs(ty: &str) -> Param {
     Param { pattern: context_pattern, typ: context_type, visibility, span: Span::default() }
 }
 
+/// Creates an initialization check to ensure that the contract has been initialized, meant to
+/// be injected as the first statement of any function after the context has been created.
+///
+/// ```noir
+/// assert_is_initialized(&mut context);
+/// ```
+fn create_init_check() -> Statement {
+    make_statement(StatementKind::Expression(call(
+        variable_path(chained_path!("aztec", "initializer", "assert_is_initialized")),
+        vec![mutable_reference("context")],
+    )))
+}
+
+/// Creates a call to mark_as_initialized which emits the initialization nullifier, meant to
+/// be injected as the last statement before returning in a constructor.
+///
+/// ```noir
+/// mark_as_initialized(&mut context);
+/// ```
+fn create_mark_as_initialized() -> Statement {
+    make_statement(StatementKind::Expression(call(
+        variable_path(chained_path!("aztec", "initializer", "mark_as_initialized")),
+        vec![mutable_reference("context")],
+    )))
+}
+
 /// Creates the private context object to be accessed within the function, the parameters need to be extracted to be
 /// appended into the args hash object.
 ///