diff --git a/noir/aztec_macros/src/lib.rs b/noir/aztec_macros/src/lib.rs index 07dca820fd9a..d08f1c81e2d9 100644 --- a/noir/aztec_macros/src/lib.rs +++ b/noir/aztec_macros/src/lib.rs @@ -63,12 +63,12 @@ impl MacroProcessor for AztecMacro { } const FUNCTION_TREE_HEIGHT: u32 = 5; -const MAX_CONTRACT_FUNCTIONS: usize = 2_usize.pow(FUNCTION_TREE_HEIGHT); +const MAX_CONTRACT_PRIVATE_FUNCTIONS: usize = 2_usize.pow(FUNCTION_TREE_HEIGHT); #[derive(Debug, Clone)] pub enum AztecMacroError { AztecDepNotFound, - ContractHasTooManyFunctions { span: Span }, + ContractHasTooManyPrivateFunctions { span: Span }, ContractConstructorMissing { span: Span }, UnsupportedFunctionArgumentType { span: Span, typ: UnresolvedTypeData }, UnsupportedStorageType { span: Option, typ: UnresolvedTypeData }, @@ -84,8 +84,8 @@ impl From for MacroError { secondary_message: None, span: None, }, - AztecMacroError::ContractHasTooManyFunctions { span } => MacroError { - primary_message: format!("Contract can only have a maximum of {} functions", MAX_CONTRACT_FUNCTIONS), + AztecMacroError::ContractHasTooManyPrivateFunctions { span } => MacroError { + primary_message: format!("Contract can only have a maximum of {} private functions", MAX_CONTRACT_PRIVATE_FUNCTIONS), secondary_message: None, span: Some(span), }, @@ -456,10 +456,22 @@ fn transform_module( if has_transformed_module { // We only want to run these checks if the macro processor has found the module to be an Aztec contract. - if module.functions.len() > MAX_CONTRACT_FUNCTIONS { + let private_functions_count = module + .functions + .iter() + .filter(|func| { + func.def + .attributes + .secondary + .iter() + .any(|attr| is_custom_attribute(attr, "aztec(private)")) + }) + .count(); + + if private_functions_count > MAX_CONTRACT_PRIVATE_FUNCTIONS { let crate_graph = &context.crate_graph[crate_id]; return Err(( - AztecMacroError::ContractHasTooManyFunctions { span: Span::default() }, + AztecMacroError::ContractHasTooManyPrivateFunctions { span: Span::default() }, crate_graph.root_file_id, )); }