refactor: cleaning up macros::functions::utils

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š <[email protected]>

Author: benesjan

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.

noir-projects/aztec-nr/aztec/src/macros/functions/external_registry.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") {

noir-projects/aztec-nr/aztec/src/macros/functions/mod.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};

noir-projects/aztec-nr/aztec/src/macros/functions/stub_registry.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;

noir-projects/aztec-nr/aztec/src/macros/functions/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")
 }