fix: Change function limit to private function limit (#4785)

AztecProtocol · Feb 27, 2024 · 2799f1f · 2799f1f
1 parent 10d90ab
commit 2799f1f
Showing 1 changed file with 18 additions and 6 deletions.
diff --git 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<Span>, typ: UnresolvedTypeData },
@@ -84,8 +84,8 @@ impl From<AztecMacroError> 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,
             ));
         }