11use crate::macros:: {
22 functions ::{
3- auth_registry::AUTHORIZE_ONCE_REGISTRY , call_interface_stubs:: stub_fn , stub_registry ,
3+ auth_registry::AUTHORIZE_ONCE_REGISTRY ,
4+ call_interface_stubs ::{
5+ register_private_fn_stub , register_public_fn_stub , register_utility_fn_stub ,
6+ },
47 },
58 notes::NOTES ,
69 utils ::{
7- fn_has_authorize_once , fn_has_noinitcheck , get_fn_visibility , is_fn_contract_library_method ,
8- is_fn_external , is_fn_initializer , is_fn_internal , is_fn_private , is_fn_test , is_fn_view ,
9- modify_fn_body , module_has_initializer , module_has_storage ,
10+ fn_has_authorize_once , fn_has_noinitcheck , is_fn_contract_library_method , is_fn_external ,
11+ is_fn_initializer , is_fn_internal , is_fn_test , is_fn_view , modify_fn_body ,
12+ module_has_initializer , module_has_storage ,
1013 },
1114};
1215use dep::protocol_types::meta::utils::derive_serialization_quotes ;
1316use std::meta:: {ctstring::AsCtString , type_of };
1417
1518pub (crate ) comptime fn transform_private (f : FunctionDefinition ) {
16- let fn_stub = stub_fn (f );
17- stub_registry:: register (f .module (), fn_stub );
19+ register_private_fn_stub (f );
1820
1921 let module_has_initializer = module_has_initializer (f .module ());
2022 let module_has_storage = module_has_storage (f .module ());
@@ -43,21 +45,29 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) {
4345 let mut context = dep::aztec::context::private_context::PrivateContext ::new (inputs , args_hash );
4446 };
4547
48+ let function_name = f .name ();
49+
4650 // Modifications introduced by the different marker attributes.
4751 let internal_check = if is_fn_internal (f ) {
48- create_internal_check (f )
52+ let assertion_message = f"Function { function_name} can only be called internally" ;
53+ quote { assert (context .msg_sender ().unwrap () == context .this_address (), $assertion_message ); }
4954 } else {
5055 quote {}
5156 };
5257
5358 let view_check = if is_fn_view (f ) {
54- create_view_check (f )
59+ let assertion_message =
60+ f"Function { function_name} can only be called statically" .as_ctstring ().as_quoted_str ();
61+ quote { assert (context .inputs .call_context .is_static_call , $assertion_message ); }
5562 } else {
5663 quote {}
5764 };
5865
5966 let (assert_initializer , mark_as_initialized ) = if is_fn_initializer (f ) {
60- (create_assert_correct_initializer_args (f ), create_mark_as_initialized (f ))
67+ (
68+ quote { aztec::macros::functions::initialization_utils:: assert_initialization_matches_address_preimage_private (context ); },
69+ quote { aztec::macros::functions::initialization_utils:: mark_as_initialized_private (&mut context ); },
70+ )
6171 } else {
6272 (quote {}, quote {})
6373 };
@@ -75,7 +85,7 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) {
7585
7686 // Initialization checks are not included in contracts that don't have initializers.
7787 let init_check = if module_has_initializer & !is_fn_initializer (f ) & !fn_has_noinitcheck (f ) {
78- create_init_check ( f )
88+ quote { aztec::macros::functions::initialization_utils:: assert_is_initialized_private (& mut context ); }
7989 } else {
8090 quote {}
8191 };
@@ -92,7 +102,7 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) {
92102
93103 // Inject the authwit check if the function is marked with #[authorize_once].
94104 let authorize_once_check = if fn_has_authorize_once (f ) {
95- create_authorize_once_check (f )
105+ create_authorize_once_check (f , true )
96106 } else {
97107 quote {}
98108 };
@@ -167,8 +177,7 @@ pub(crate) comptime fn transform_private(f: FunctionDefinition) {
167177}
168178
169179pub (crate ) comptime fn transform_public (f : FunctionDefinition ) {
170- let fn_stub = stub_fn (f );
171- stub_registry:: register (f .module (), fn_stub );
180+ register_public_fn_stub (f );
172181
173182 let module_has_initializer = module_has_initializer (f .module ());
174183 let module_has_storage = module_has_storage (f .module ());
@@ -199,21 +208,29 @@ pub(crate) comptime fn transform_public(f: FunctionDefinition) {
199208 });
200209 };
201210
211+ let name = f .name ();
202212 // Modifications introduced by the different marker attributes.
203213 let internal_check = if is_fn_internal (f ) {
204- create_internal_check (f )
214+ let assertion_message = f"Function { name} can only be called internally" ;
215+ quote { assert (context .msg_sender ().unwrap () == context .this_address (), $assertion_message ); }
205216 } else {
206217 quote {}
207218 };
208219
209220 let view_check = if is_fn_view (f ) {
210- create_view_check (f )
221+ let name = f .name ();
222+ let assertion_message =
223+ f"Function { name} can only be called statically" .as_ctstring ().as_quoted_str ();
224+ quote { assert (context .is_static_call (), $assertion_message ); }
211225 } else {
212226 quote {}
213227 };
214228
215229 let (assert_initializer , mark_as_initialized ) = if is_fn_initializer (f ) {
216- (create_assert_correct_initializer_args (f ), create_mark_as_initialized (f ))
230+ (
231+ quote { aztec::macros::functions::initialization_utils:: assert_initialization_matches_address_preimage_public (context ); },
232+ quote { aztec::macros::functions::initialization_utils:: mark_as_initialized_public (&mut context ); },
233+ )
217234 } else {
218235 (quote {}, quote {})
219236 };
@@ -231,14 +248,14 @@ pub(crate) comptime fn transform_public(f: FunctionDefinition) {
231248
232249 // Initialization checks are not included in contracts that don't have initializers.
233250 let init_check = if module_has_initializer & !fn_has_noinitcheck (f ) & !is_fn_initializer (f ) {
234- create_init_check ( f )
251+ quote { aztec::macros::functions::initialization_utils:: assert_is_initialized_public (& mut context ); }
235252 } else {
236253 quote {}
237254 };
238255
239256 // Inject the authwit check if the function is marked with #[authorize_once].
240257 let authorize_once_check = if fn_has_authorize_once (f ) {
241- create_authorize_once_check (f )
258+ create_authorize_once_check (f , false )
242259 } else {
243260 quote {}
244261 };
@@ -269,8 +286,7 @@ pub(crate) comptime fn transform_public(f: FunctionDefinition) {
269286}
270287
271288pub (crate ) comptime fn transform_utility (f : FunctionDefinition ) {
272- let fn_stub = stub_fn (f );
273- stub_registry:: register (f .module (), fn_stub );
289+ register_utility_fn_stub (f );
274290
275291 // Create utility context
276292 let context_creation =
@@ -307,42 +323,6 @@ pub(crate) comptime fn transform_utility(f: FunctionDefinition) {
307323 f .set_return_public (true );
308324}
309325
310- comptime fn create_internal_check (f : FunctionDefinition ) -> Quoted {
311- let name = f .name ();
312- let assertion_message = f"Function { name} can only be called internally" ;
313- quote { assert (context .msg_sender ().unwrap () == context .this_address (), $assertion_message ); }
314- }
315-
316- comptime fn create_view_check (f : FunctionDefinition ) -> Quoted {
317- let name = f .name ();
318- let assertion_message = f"Function { name} can only be called statically" ;
319- if is_fn_private (f ) {
320- // Here `context` is of type context::PrivateContext
321- quote { assert (context .inputs .call_context .is_static_call , $assertion_message ); }
322- } else {
323- // Here `context` is of type context::PublicContext
324- quote { assert (context .is_static_call (), $assertion_message ); }
325- }
326- }
327-
328- comptime fn create_assert_correct_initializer_args (f : FunctionDefinition ) -> Quoted {
329- let fn_visibility = get_fn_visibility (f );
330- f"dep::aztec::macros::functions::initialization_utils::assert_initialization_matches_address_preimage_{ fn_visibility} (context);"
331- .quoted_contents ()
332- }
333-
334- comptime fn create_mark_as_initialized (f : FunctionDefinition ) -> Quoted {
335- let fn_visibility = get_fn_visibility (f );
336- f"dep::aztec::macros::functions::initialization_utils::mark_as_initialized_{ fn_visibility} (&mut context);"
337- .quoted_contents ()
338- }
339-
340- comptime fn create_init_check (f : FunctionDefinition ) -> Quoted {
341- let fn_visibility = get_fn_visibility (f );
342- f"dep::aztec::macros::functions::initialization_utils::assert_is_initialized_{ fn_visibility} (&mut context);"
343- .quoted_contents ()
344- }
345-
346326/// Injects a call to `aztec::messages::discovery::discover_new_messages`, causing for new notes to be added to PXE and made
347327/// available for the current execution.
348328pub (crate ) comptime fn create_message_discovery_call () -> Quoted {
@@ -369,7 +349,17 @@ pub(crate) comptime fn create_message_discovery_call() -> Quoted {
369349/// where `from` and `authwit_nonce` are the names of the parameters that are expected to be present in the function definition.
370350/// This check is injected by the `#[authorize_once("from_arg_name", "nonce_arg_name")]`, which allows the user to define
371351/// which parameters to use.
372- pub (crate ) comptime fn create_authorize_once_check (f : FunctionDefinition ) -> Quoted {
352+ ///
353+ /// # Arguments
354+ /// * `f` - The function definition to inject the authwit verification check into. The function must have parameters
355+ /// matching the names specified in the `#[authorize_once]` attribute.
356+ /// * `is_private` - Whether the function is a private function (`true`) or a public function (`false`). This determines
357+ /// which authwit verification method to use: `assert_current_call_valid_authwit` for private functions
358+ /// or `assert_current_call_valid_authwit_public` for public functions.
359+ pub (crate ) comptime fn create_authorize_once_check (
360+ f : FunctionDefinition ,
361+ is_private : bool ,
362+ ) -> Quoted {
373363 let maybe_authorize_once_args = AUTHORIZE_ONCE_REGISTRY .get (f );
374364 let authorize_once_args = if maybe_authorize_once_args .is_some () {
375365 maybe_authorize_once_args .unwrap ()
@@ -417,7 +407,7 @@ pub(crate) comptime fn create_authorize_once_check(f: FunctionDefinition) -> Quo
417407
418408 let nonce_check_quote = f"{ nonce_arg_name_quoted} == 0" .quoted_contents ();
419409
420- let fn_call = if is_fn_private ( f ) {
410+ let fn_call = if is_private {
421411 // At this point, the original args of the fn have already been altered by the macro
422412 // to include PrivateContextInputs, so we need to adjust the args_len accordingly.
423413 let args_len = f .parameters ().len () - 1 ;
0 commit comments