Split aztec attributes

AztecProtocol · Feb 28, 2024 · ebc3832 · ebc3832
1 parent f826451
commit ebc3832
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 52 deletions.
diff --git a/noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr b/noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr
@@ -17,13 +17,15 @@ contract StatefulTest {
         public_values: Map<AztecAddress, PublicMutable<Field>>,
     }
 
-    #[aztec(private,initializer)]
+    #[aztec(private)]
+    #[aztec(initializer)]
     fn constructor(owner: AztecAddress, value: Field) {
         let selector = FunctionSelector::from_signature("internal_create_note((Field),Field)");
         let _res = context.call_private_function(context.this_address(), selector, [owner.to_field(), value]);
     }
 
-    #[aztec(private,initcheck)]
+    #[aztec(private)]
+    #[aztec(initcheck)]
     fn create_note(owner: AztecAddress, value: Field) {
         if (value != 0) {
             let loc = storage.notes.at(owner);

diff --git a/noir/aztec_macros/src/lib.rs b/noir/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,62 +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)") {
-                let is_initializer = false;
-                let skip_init_check = true;
-                transform_function(
-                    "Private",
-                    func,
-                    storage_defined,
-                    is_initializer,
-                    skip_init_check,
-                )
-                .map_err(|err| (err, crate_graph.root_file_id))?;
-                has_transformed_module = true;
-            } else if is_custom_attribute(&secondary_attribute, "aztec(private,initializer)") {
-                let is_initializer = true;
-                let skip_init_check = false;
-                transform_function(
-                    "Private",
-                    func,
-                    storage_defined,
-                    is_initializer,
-                    skip_init_check,
-                )
-                .map_err(|err| (err, crate_graph.root_file_id))?;
-                has_transformed_module = true;
-            } else if is_custom_attribute(&secondary_attribute, "aztec(private,initcheck)") {
-                let is_initializer = false;
-                let skip_init_check = false;
-                transform_function(
-                    "Private",
-                    func,
-                    storage_defined,
-                    is_initializer,
-                    skip_init_check,
-                )
-                .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)") {
-                let is_initializer = false;
-                let skip_init_check = true;
-                transform_function(
-                    "Public",
-                    func,
-                    storage_defined,
-                    is_initializer,
-                    skip_init_check,
-                )
-                .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);
@@ -676,6 +663,15 @@ fn transform_function(
 
     // 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);
     }
@@ -705,6 +701,15 @@ fn transform_function(
 
     // 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);
     }