From 5250359944ac50310c0e7e82b4d5f7bf11a8a89e Mon Sep 17 00:00:00 2001 From: benesjan <13470840+benesjan@users.noreply.github.com> Date: Fri, 24 Oct 2025 01:11:59 +0000 Subject: [PATCH] refactor: cleaning up macros::functions::utils MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR was quite a journey: ## Step 1 When trying to figure out how to rename the incorrectly named `get_fn_visibility` in a PR down the stack I realized that it's better to just drop it because: - it was only used by `create_assert_correct_initializer_args`, `create_mark_as_initialized` and `create_init_check` functions which in turn where only called in `transfrom_private` and `transform_public`, - these functions contained only 2 lines of code. Because of this ^ it's better to just directly inline contents of these funcs into the `transform_*` functions. When I finished this I checked where are the is `is_fn_private` and `is_fn_public` utils used which led me to step 2: ## Step 2 When checking when are the above mentioned functions used I realized that we use them in `create_authorize_once_check` and in the `stub_fn`. In both of these cases the code was just terrible because these functions are called from `transform_private` and `transform_public` functions hence the flow was: `transform_private` --> `stub_fn` --> are we in private? --> if yes call `create_private_stub` which I could replace with: `transform_private` --> `create_private_stub` So I dropped the `stub_fn` and just called the `create_private_stub` directly. The similar was the case for `create_authorize_once_check` where I could simply pass in an `is_private` arg to the function and get rid of the `is_fn_private` dependency. ## Step 3 Once I did this I realized that I don't need the `EXTERNAL_REGISTRY` I introduced in the PR down the stack at all, removing a lot of the sad ugly complexity I needed to introduce there 🎉 ## Future work - The whole `noir-projects/aztec-nr/aztec/src/macros/functions` directory is very messy and I plan on move the code around into more logical places in the dir once this stack of PRs is merged. - I think the `stub` name is non-descriptive enough so will try to come up with something better. Co-authored-by: Jan Beneš --- .../macros/functions/call_interface_stubs.nr | 50 +++++---- .../src/macros/functions/external_registry.nr | 24 ---- .../aztec/src/macros/functions/mod.nr | 4 - .../src/macros/functions/stub_registry.nr | 1 + .../aztec/src/macros/functions/utils.nr | 106 ++++++++---------- .../aztec-nr/aztec/src/macros/utils.nr | 35 ------ 6 files changed, 76 insertions(+), 144 deletions(-) delete mode 100644 noir-projects/aztec-nr/aztec/src/macros/functions/external_registry.nr diff --git a/noir-projects/aztec-nr/aztec/src/macros/functions/call_interface_stubs.nr b/noir-projects/aztec-nr/aztec/src/macros/functions/call_interface_stubs.nr index 8bf0b63a100e..b2bf7bb1f55f 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/functions/call_interface_stubs.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/functions/call_interface_stubs.nr @@ -1,10 +1,11 @@ -// This file generates call interface stub functions for Aztec contract methods. -// It creates wrapper functions that serialize arguments and return appropriate call interfaces -// for invoking contract functions across different contexts (private, public, static, utility). - -use crate::macros::utils::{ - AsStrQuote, compute_fn_selector, is_fn_private, is_fn_public, is_fn_view, -}; +//! Stubs are auto-generated wrapper functions that provide an ergonomic interface for cross-contract calls. +//! Instead of manually serializing arguments and creating call interfaces, stubs allow natural syntax, e.g. for +//! enqueuing calls to public functions: +//! +//! ExternalContract.at(address).some_method(arg1, arg2).enqueue() + +use crate::macros::utils::{AsStrQuote, compute_fn_selector, is_fn_view}; +use super::stub_registry; use protocol_types::meta::utils::derive_serialization_quotes; use std::meta::unquote; @@ -18,24 +19,27 @@ comptime global FROM_FIELD: TypedExpr = { .as_typed_expr() }; -pub comptime fn stub_fn(f: FunctionDefinition) -> Quoted { - let is_static_call = is_fn_view(f); +pub comptime fn register_private_fn_stub(f: FunctionDefinition) { + let stub = if is_fn_view(f) { + create_private_static_stub(f) + } else { + create_private_stub(f) + }; + stub_registry::register(f.module(), stub); +} - if is_fn_private(f) { - if is_static_call { - create_private_static_stub(f) - } else { - create_private_stub(f) - } - } else if is_fn_public(f) { - if is_static_call { - create_public_static_stub(f) - } else { - create_public_stub(f) - } +pub comptime fn register_public_fn_stub(f: FunctionDefinition) { + let stub = if is_fn_view(f) { + create_public_static_stub(f) } else { - create_utility_stub(f) - } + create_public_stub(f) + }; + stub_registry::register(f.module(), stub); +} + +pub comptime fn register_utility_fn_stub(f: FunctionDefinition) { + let stub = create_utility_stub(f); + stub_registry::register(f.module(), stub); } /// Utility function creating stubs used by all the stub functions in this file. diff --git a/noir-projects/aztec-nr/aztec/src/macros/functions/external_registry.nr b/noir-projects/aztec-nr/aztec/src/macros/functions/external_registry.nr deleted file mode 100644 index 4bf1e170570a..000000000000 --- a/noir-projects/aztec-nr/aztec/src/macros/functions/external_registry.nr +++ /dev/null @@ -1,24 +0,0 @@ -use poseidon::poseidon2::Poseidon2Hasher; -use std::{collections::umap::UHashMap, hash::BuildHasherDefault}; - -/// Key is external function definition and value is the argument of the #[external] attribute ("private", "public", -/// or "utility"). -comptime mut global EXTERNAL_REGISTRY: UHashMap> = - UHashMap::default(); - -/// Registers functions in the EXTERNAL_REGISTRY map. Called by the #[external] macro. -pub(crate) comptime fn register(f: FunctionDefinition, f_type: CtString) { - // Check if function is already registered as external - if EXTERNAL_REGISTRY.get(f).is_some() { - let name = f.name(); - panic( - f"The #[external] attribute cannot be applied multiple times to the same function - {name} already has it", - ); - } - - EXTERNAL_REGISTRY.insert(f, f_type); -} - -pub(crate) comptime fn get(f: FunctionDefinition) -> Option { - EXTERNAL_REGISTRY.get(f) -} diff --git a/noir-projects/aztec-nr/aztec/src/macros/functions/mod.nr b/noir-projects/aztec-nr/aztec/src/macros/functions/mod.nr index 5529ef10bc42..1fc8c0509ba4 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/functions/mod.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/functions/mod.nr @@ -5,7 +5,6 @@ pub(crate) mod call_interface_stubs; pub mod initialization_utils; pub(crate) mod stub_registry; pub(crate) mod auth_registry; -pub(crate) mod external_registry; pub(crate) mod utils; use crate::macros::{ @@ -134,9 +133,6 @@ pub comptime fn authorize_once( /// Same as in Solidity external functions are functions that our callable from outside the contract. External /// functions can be either private, public, or utility. pub comptime fn external(f: FunctionDefinition, f_type: CtString) -> Quoted { - // Register the function type in the global registry so helper functions can check it - external_registry::register(f, f_type); - if f_type.eq("private") { private(f) } else if f_type.eq("public") { diff --git a/noir-projects/aztec-nr/aztec/src/macros/functions/stub_registry.nr b/noir-projects/aztec-nr/aztec/src/macros/functions/stub_registry.nr index 68207228f861..88fa3f551db6 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/functions/stub_registry.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/functions/stub_registry.nr @@ -1,6 +1,7 @@ // This file provides a registry for storing and retrieving function stubs generated by macros. // It maintains a global hashmap that maps modules to their generated stub functions, // allowing the macro system to collect and organize contract interface stubs. +// See the call_interface_stubs module for more details on what stubs are and how they are generated. use poseidon::poseidon2::Poseidon2Hasher; use std::{collections::umap::UHashMap, hash::BuildHasherDefault}; diff --git a/noir-projects/aztec-nr/aztec/src/macros/functions/utils.nr b/noir-projects/aztec-nr/aztec/src/macros/functions/utils.nr index cc8c0333473a..540c30f1f4f8 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/functions/utils.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/functions/utils.nr @@ -1,20 +1,22 @@ use crate::macros::{ functions::{ - auth_registry::AUTHORIZE_ONCE_REGISTRY, call_interface_stubs::stub_fn, stub_registry, + auth_registry::AUTHORIZE_ONCE_REGISTRY, + call_interface_stubs::{ + register_private_fn_stub, register_public_fn_stub, register_utility_fn_stub, + }, }, notes::NOTES, utils::{ - fn_has_authorize_once, fn_has_noinitcheck, get_fn_visibility, is_fn_contract_library_method, - is_fn_external, is_fn_initializer, is_fn_internal, is_fn_private, is_fn_test, is_fn_view, - modify_fn_body, module_has_initializer, module_has_storage, + fn_has_authorize_once, fn_has_noinitcheck, is_fn_contract_library_method, is_fn_external, + is_fn_initializer, is_fn_internal, is_fn_test, is_fn_view, modify_fn_body, + module_has_initializer, module_has_storage, }, }; use dep::protocol_types::meta::utils::derive_serialization_quotes; use std::meta::{ctstring::AsCtString, type_of}; pub(crate) comptime fn transform_private(f: FunctionDefinition) { - let fn_stub = stub_fn(f); - stub_registry::register(f.module(), fn_stub); + register_private_fn_stub(f); let module_has_initializer = module_has_initializer(f.module()); let module_has_storage = module_has_storage(f.module()); @@ -43,21 +45,29 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) { let mut context = dep::aztec::context::private_context::PrivateContext::new(inputs, args_hash); }; + let function_name = f.name(); + // Modifications introduced by the different marker attributes. let internal_check = if is_fn_internal(f) { - create_internal_check(f) + let assertion_message = f"Function {function_name} can only be called internally"; + quote { assert(context.msg_sender().unwrap() == context.this_address(), $assertion_message); } } else { quote {} }; let view_check = if is_fn_view(f) { - create_view_check(f) + let assertion_message = + f"Function {function_name} can only be called statically".as_ctstring().as_quoted_str(); + quote { assert(context.inputs.call_context.is_static_call, $assertion_message); } } else { quote {} }; let (assert_initializer, mark_as_initialized) = if is_fn_initializer(f) { - (create_assert_correct_initializer_args(f), create_mark_as_initialized(f)) + ( + quote { aztec::macros::functions::initialization_utils::assert_initialization_matches_address_preimage_private(context); }, + quote { aztec::macros::functions::initialization_utils::mark_as_initialized_private(&mut context); }, + ) } else { (quote {}, quote {}) }; @@ -75,7 +85,7 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) { // Initialization checks are not included in contracts that don't have initializers. let init_check = if module_has_initializer & !is_fn_initializer(f) & !fn_has_noinitcheck(f) { - create_init_check(f) + quote { aztec::macros::functions::initialization_utils::assert_is_initialized_private(&mut context); } } else { quote {} }; @@ -92,7 +102,7 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) { // Inject the authwit check if the function is marked with #[authorize_once]. let authorize_once_check = if fn_has_authorize_once(f) { - create_authorize_once_check(f) + create_authorize_once_check(f, true) } else { quote {} }; @@ -167,8 +177,7 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) { } pub(crate) comptime fn transform_public(f: FunctionDefinition) { - let fn_stub = stub_fn(f); - stub_registry::register(f.module(), fn_stub); + register_public_fn_stub(f); let module_has_initializer = module_has_initializer(f.module()); let module_has_storage = module_has_storage(f.module()); @@ -199,21 +208,29 @@ pub(crate) comptime fn transform_public(f: FunctionDefinition) { }); }; + let name = f.name(); // Modifications introduced by the different marker attributes. let internal_check = if is_fn_internal(f) { - create_internal_check(f) + let assertion_message = f"Function {name} can only be called internally"; + quote { assert(context.msg_sender().unwrap() == context.this_address(), $assertion_message); } } else { quote {} }; let view_check = if is_fn_view(f) { - create_view_check(f) + let name = f.name(); + let assertion_message = + f"Function {name} can only be called statically".as_ctstring().as_quoted_str(); + quote { assert(context.is_static_call(), $assertion_message); } } else { quote {} }; let (assert_initializer, mark_as_initialized) = if is_fn_initializer(f) { - (create_assert_correct_initializer_args(f), create_mark_as_initialized(f)) + ( + quote { aztec::macros::functions::initialization_utils::assert_initialization_matches_address_preimage_public(context); }, + quote { aztec::macros::functions::initialization_utils::mark_as_initialized_public(&mut context); }, + ) } else { (quote {}, quote {}) }; @@ -231,14 +248,14 @@ pub(crate) comptime fn transform_public(f: FunctionDefinition) { // Initialization checks are not included in contracts that don't have initializers. let init_check = if module_has_initializer & !fn_has_noinitcheck(f) & !is_fn_initializer(f) { - create_init_check(f) + quote { aztec::macros::functions::initialization_utils::assert_is_initialized_public(&mut context); } } else { quote {} }; // Inject the authwit check if the function is marked with #[authorize_once]. let authorize_once_check = if fn_has_authorize_once(f) { - create_authorize_once_check(f) + create_authorize_once_check(f, false) } else { quote {} }; @@ -269,8 +286,7 @@ pub(crate) comptime fn transform_public(f: FunctionDefinition) { } pub(crate) comptime fn transform_utility(f: FunctionDefinition) { - let fn_stub = stub_fn(f); - stub_registry::register(f.module(), fn_stub); + register_utility_fn_stub(f); // Create utility context let context_creation = @@ -307,42 +323,6 @@ pub(crate) comptime fn transform_utility(f: FunctionDefinition) { f.set_return_public(true); } -comptime fn create_internal_check(f: FunctionDefinition) -> Quoted { - let name = f.name(); - let assertion_message = f"Function {name} can only be called internally"; - quote { assert(context.msg_sender().unwrap() == context.this_address(), $assertion_message); } -} - -comptime fn create_view_check(f: FunctionDefinition) -> Quoted { - let name = f.name(); - let assertion_message = f"Function {name} can only be called statically"; - if is_fn_private(f) { - // Here `context` is of type context::PrivateContext - quote { assert(context.inputs.call_context.is_static_call, $assertion_message); } - } else { - // Here `context` is of type context::PublicContext - quote { assert(context.is_static_call(), $assertion_message); } - } -} - -comptime fn create_assert_correct_initializer_args(f: FunctionDefinition) -> Quoted { - let fn_visibility = get_fn_visibility(f); - f"dep::aztec::macros::functions::initialization_utils::assert_initialization_matches_address_preimage_{fn_visibility}(context);" - .quoted_contents() -} - -comptime fn create_mark_as_initialized(f: FunctionDefinition) -> Quoted { - let fn_visibility = get_fn_visibility(f); - f"dep::aztec::macros::functions::initialization_utils::mark_as_initialized_{fn_visibility}(&mut context);" - .quoted_contents() -} - -comptime fn create_init_check(f: FunctionDefinition) -> Quoted { - let fn_visibility = get_fn_visibility(f); - f"dep::aztec::macros::functions::initialization_utils::assert_is_initialized_{fn_visibility}(&mut context);" - .quoted_contents() -} - /// Injects a call to `aztec::messages::discovery::discover_new_messages`, causing for new notes to be added to PXE and made /// available for the current execution. pub(crate) comptime fn create_message_discovery_call() -> Quoted { @@ -369,7 +349,17 @@ pub(crate) comptime fn create_message_discovery_call() -> Quoted { /// where `from` and `authwit_nonce` are the names of the parameters that are expected to be present in the function definition. /// This check is injected by the `#[authorize_once("from_arg_name", "nonce_arg_name")]`, which allows the user to define /// which parameters to use. -pub(crate) comptime fn create_authorize_once_check(f: FunctionDefinition) -> Quoted { +/// +/// # Arguments +/// * `f` - The function definition to inject the authwit verification check into. The function must have parameters +/// matching the names specified in the `#[authorize_once]` attribute. +/// * `is_private` - Whether the function is a private function (`true`) or a public function (`false`). This determines +/// which authwit verification method to use: `assert_current_call_valid_authwit` for private functions +/// or `assert_current_call_valid_authwit_public` for public functions. +pub(crate) comptime fn create_authorize_once_check( + f: FunctionDefinition, + is_private: bool, +) -> Quoted { let maybe_authorize_once_args = AUTHORIZE_ONCE_REGISTRY.get(f); let authorize_once_args = if maybe_authorize_once_args.is_some() { maybe_authorize_once_args.unwrap() @@ -417,7 +407,7 @@ pub(crate) comptime fn create_authorize_once_check(f: FunctionDefinition) -> Quo let nonce_check_quote = f"{nonce_arg_name_quoted} == 0".quoted_contents(); - let fn_call = if is_fn_private(f) { + let fn_call = if is_private { // At this point, the original args of the fn have already been altered by the macro // to include PrivateContextInputs, so we need to adjust the args_len accordingly. let args_len = f.parameters().len() - 1; diff --git a/noir-projects/aztec-nr/aztec/src/macros/utils.nr b/noir-projects/aztec-nr/aztec/src/macros/utils.nr index a23f2e50c3e4..42f34150b7da 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/utils.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/utils.nr @@ -1,45 +1,10 @@ -use crate::macros::functions::external_registry; use dep::protocol_types::meta::derive_serialize; use std::meta::unquote; -pub(crate) comptime fn get_fn_visibility(f: FunctionDefinition) -> Quoted { - let maybe_external_type = external_registry::get(f); - if maybe_external_type.is_some() { - let external_type = maybe_external_type.unwrap(); - if external_type.eq("private") { - quote { private } - } else if external_type.eq("public") { - quote { public } - } else { - panic(f"Function is neither private nor public") - } - } else { - panic(f"Function is neither private nor public") - } -} - pub(crate) comptime fn is_fn_external(f: FunctionDefinition) -> bool { f.has_named_attribute("external") } -pub(crate) comptime fn is_fn_private(f: FunctionDefinition) -> bool { - let maybe_external_type = external_registry::get(f); - if maybe_external_type.is_some() { - maybe_external_type.unwrap_unchecked().eq("private") - } else { - false - } -} - -pub(crate) comptime fn is_fn_public(f: FunctionDefinition) -> bool { - let maybe_external_type = external_registry::get(f); - if maybe_external_type.is_some() { - maybe_external_type.unwrap_unchecked().eq("public") - } else { - false - } -} - pub(crate) comptime fn is_fn_contract_library_method(f: FunctionDefinition) -> bool { f.has_named_attribute("contract_library_method") }