diff --git a/noir-projects/aztec-nr/aztec/src/macros/dispatch.nr b/noir-projects/aztec-nr/aztec/src/macros/dispatch.nr index 61f0f50eef23..40536b8f999d 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/dispatch.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/dispatch.nr @@ -1,4 +1,4 @@ -use super::utils::compute_fn_selector; +use super::utils::{compute_fn_selector, size_in_fields}; use poseidon::poseidon2::Poseidon2Hasher; use std::{collections::umap::UHashMap, hash::BuildHasherDefault, panic}; @@ -119,83 +119,6 @@ pub comptime fn generate_public_dispatch(m: Module) -> Quoted { } } -comptime fn size_in_fields(typ: Type) -> u32 { - let size = array_size_in_fields(typ); - let size = size.or_else(|| bool_size_in_fields(typ)); - let size = size.or_else(|| constant_size_in_fields(typ)); - let size = size.or_else(|| field_size_in_fields(typ)); - let size = size.or_else(|| int_size_in_fields(typ)); - let size = size.or_else(|| str_size_in_fields(typ)); - let size = size.or_else(|| struct_size_in_fields(typ)); - let size = size.or_else(|| tuple_size_in_fields(typ)); - if size.is_some() { - size.unwrap() - } else { - panic(f"Can't determine size in fields of {typ}") - } -} - -comptime fn array_size_in_fields(typ: Type) -> Option { - typ.as_array().and_then(|typ: (Type, Type)| { - let (typ, element_size) = typ; - element_size.as_constant().map(|x: u32| x * size_in_fields(typ)) - }) -} - -comptime fn bool_size_in_fields(typ: Type) -> Option { - if typ.is_bool() { - Option::some(1) - } else { - Option::none() - } -} - -comptime fn field_size_in_fields(typ: Type) -> Option { - if typ.is_field() { - Option::some(1) - } else { - Option::none() - } -} - -comptime fn int_size_in_fields(typ: Type) -> Option { - if typ.as_integer().is_some() { - Option::some(1) - } else { - Option::none() - } -} - -comptime fn constant_size_in_fields(typ: Type) -> Option { - typ.as_constant() -} - -comptime fn str_size_in_fields(typ: Type) -> Option { - typ.as_str().map(|typ| size_in_fields(typ)) -} - -comptime fn struct_size_in_fields(typ: Type) -> Option { - typ.as_data_type().map(|typ: (TypeDefinition, [Type])| { - let struct_type = typ.0; - let generics = typ.1; - let mut size = 0; - for field in struct_type.fields(generics) { - size += size_in_fields(field.1); - } - size - }) -} - -comptime fn tuple_size_in_fields(typ: Type) -> Option { - typ.as_tuple().map(|types: [Type]| { - let mut size = 0; - for typ in types { - size += size_in_fields(typ); - } - size - }) -} - comptime fn get_type() -> Type { let t: T = std::mem::zeroed(); std::meta::type_of(t) 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 b0b9cbe4cd37..0909e93491d1 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,5 +1,6 @@ use crate::macros::utils::{ add_to_field_slice, AsStrQuote, compute_fn_selector, is_fn_private, is_fn_public, is_fn_view, + size_in_fields, }; use std::meta::{type_of, unquote}; @@ -100,9 +101,10 @@ comptime fn create_private_stub(f: FunctionDefinition) -> Quoted { let (fn_name, fn_parameters_list, serialized_args_slice_construction, fn_name_str, fn_name_len, fn_selector) = create_stub_base(f); let fn_return_type = f.return_type(); + let size = size_in_fields(fn_return_type); quote { - pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PrivateCallInterface<$fn_name_len, $fn_return_type, _> { + pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PrivateCallInterface<$fn_name_len, $fn_return_type, $size> { $serialized_args_slice_construction let selector = $FROM_FIELD($fn_selector); dep::aztec::context::call_interfaces::PrivateCallInterface::new( @@ -120,9 +122,10 @@ comptime fn create_private_static_stub(f: FunctionDefinition) -> Quoted { let (fn_name, fn_parameters_list, serialized_args_slice_construction, fn_name_str, fn_name_len, fn_selector) = create_stub_base(f); let fn_return_type = f.return_type(); + let size = size_in_fields(fn_return_type); quote { - pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PrivateStaticCallInterface<$fn_name_len, $fn_return_type, _> { + pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PrivateStaticCallInterface<$fn_name_len, $fn_return_type, $size> { $serialized_args_slice_construction let selector = $FROM_FIELD($fn_selector); dep::aztec::context::call_interfaces::PrivateStaticCallInterface::new( @@ -176,9 +179,10 @@ comptime fn create_public_stub(f: FunctionDefinition) -> Quoted { let (fn_name, fn_parameters_list, serialized_args_slice_construction, fn_name_str, fn_name_len, fn_selector) = create_stub_base(f); let fn_return_type = f.return_type(); + let size = size_in_fields(fn_return_type); quote { - pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PublicCallInterface<$fn_name_len, $fn_return_type, _> { + pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PublicCallInterface<$fn_name_len, $fn_return_type, $size> { $serialized_args_slice_construction let selector = $FROM_FIELD($fn_selector); dep::aztec::context::call_interfaces::PublicCallInterface::new( @@ -196,9 +200,10 @@ comptime fn create_public_static_stub(f: FunctionDefinition) -> Quoted { let (fn_name, fn_parameters_list, serialized_args_slice_construction, fn_name_str, fn_name_len, fn_selector) = create_stub_base(f); let fn_return_type = f.return_type(); + let size = size_in_fields(fn_return_type); quote { - pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PublicStaticCallInterface<$fn_name_len, $fn_return_type, _> { + pub fn $fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::PublicStaticCallInterface<$fn_name_len, $fn_return_type, $size> { $serialized_args_slice_construction let selector = $FROM_FIELD($fn_selector); dep::aztec::context::call_interfaces::PublicStaticCallInterface::new( @@ -252,12 +257,13 @@ comptime fn create_utility_stub(f: FunctionDefinition) -> Quoted { let (fn_name, fn_parameters_list, serialized_args_slice_construction, fn_name_str, fn_name_len, fn_selector) = create_stub_base(f); let fn_return_type = f.return_type(); + let size = size_in_fields(fn_return_type); // This is here because utility function call interfaces can only be used within TXe tests. let modified_fn_name = f"_experimental_{fn_name}".quoted_contents(); quote { - pub fn $modified_fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::UtilityCallInterface<$fn_name_len, $fn_return_type, _> { + pub fn $modified_fn_name(self, $fn_parameters_list) -> dep::aztec::context::call_interfaces::UtilityCallInterface<$fn_name_len, $fn_return_type, $size> { $serialized_args_slice_construction let selector = $FROM_FIELD($fn_selector); dep::aztec::context::call_interfaces::UtilityCallInterface::new( diff --git a/noir-projects/aztec-nr/aztec/src/macros/utils.nr b/noir-projects/aztec-nr/aztec/src/macros/utils.nr index a8b0383aadc5..75da12ac9494 100644 --- a/noir-projects/aztec-nr/aztec/src/macros/utils.nr +++ b/noir-projects/aztec-nr/aztec/src/macros/utils.nr @@ -288,3 +288,76 @@ pub(crate) comptime fn get_trait_impl_method( .filter(|m| m.name() == target_method)[0] .as_typed_expr() } + +pub comptime fn size_in_fields(typ: Type) -> u32 { + array_size_in_fields(typ) + .or_else(|| bool_size_in_fields(typ)) + .or_else(|| constant_size_in_fields(typ)) + .or_else(|| field_size_in_fields(typ)) + .or_else(|| int_size_in_fields(typ)) + .or_else(|| str_size_in_fields(typ)) + .or_else(|| struct_size_in_fields(typ)) + .or_else(|| tuple_size_in_fields(typ)) + .expect(f"Can't determine size in fields of {typ}") +} + +comptime fn array_size_in_fields(typ: Type) -> Option { + typ.as_array().and_then(|typ: (Type, Type)| { + let (typ, element_size) = typ; + element_size.as_constant().map(|x: u32| x * size_in_fields(typ)) + }) +} + +comptime fn bool_size_in_fields(typ: Type) -> Option { + if typ.is_bool() { + Option::some(1) + } else { + Option::none() + } +} + +comptime fn field_size_in_fields(typ: Type) -> Option { + if typ.is_field() { + Option::some(1) + } else { + Option::none() + } +} + +comptime fn int_size_in_fields(typ: Type) -> Option { + if typ.as_integer().is_some() { + Option::some(1) + } else { + Option::none() + } +} + +comptime fn constant_size_in_fields(typ: Type) -> Option { + typ.as_constant() +} + +comptime fn str_size_in_fields(typ: Type) -> Option { + typ.as_str().map(|typ| size_in_fields(typ)) +} + +comptime fn struct_size_in_fields(typ: Type) -> Option { + typ.as_data_type().map(|typ: (TypeDefinition, [Type])| { + let struct_type = typ.0; + let generics = typ.1; + let mut size = 0; + for field in struct_type.fields(generics) { + size += size_in_fields(field.1); + } + size + }) +} + +comptime fn tuple_size_in_fields(typ: Type) -> Option { + typ.as_tuple().map(|types: [Type]| { + let mut size = 0; + for typ in types { + size += size_in_fields(typ); + } + size + }) +}