From dc3138ba8313c5ac0c709b8f0524d6751b3d729b Mon Sep 17 00:00:00 2001 From: Santiago Palladino Date: Wed, 28 Feb 2024 10:06:04 -0300 Subject: [PATCH] feat: Macros for initializer checks --- .../stateful_test_contract/src/main.nr | 4 +- noir/noir-repo/aztec_macros/src/lib.rs | 105 ++++++++++++++++-- 2 files changed, 97 insertions(+), 12 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 index 88b60094c83..d0738221169 100644 --- a/noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr +++ b/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); diff --git a/noir/noir-repo/aztec_macros/src/lib.rs b/noir/noir-repo/aztec_macros/src/lib.rs index 2e4637b7732..2516b380ff6 100644 --- a/noir/noir-repo/aztec_macros/src/lib.rs +++ b/noir/noir-repo/aztec_macros/src/lib.rs @@ -74,6 +74,7 @@ pub enum AztecMacroError { UnsupportedStorageType { span: Option, typ: UnresolvedTypeData }, CouldNotAssignStorageSlots { secondary_message: Option }, EventError { span: Span, message: String }, + UnsupportedAttributes { span: Span, secondary_message: Option }, } impl From for MacroError { @@ -114,6 +115,11 @@ impl From 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. ///