From f21b39694f47f40cdc2c17ea264c8e65ca38418c Mon Sep 17 00:00:00 2001
From: Tom French <tom@tomfren.ch>
Date: Fri, 12 Jan 2024 16:48:52 +0000
Subject: [PATCH 1/2] feat: add `noir-compiler` checks to `aztec_macros`

---
 aztec_macros/src/lib.rs | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/aztec_macros/src/lib.rs b/aztec_macros/src/lib.rs
index 24f0f5a71b6..4bc96e3bfeb 100644
--- a/aztec_macros/src/lib.rs
+++ b/aztec_macros/src/lib.rs
@@ -31,10 +31,15 @@ impl MacroProcessor for AztecMacro {
     }
 }
 
+const FUNCTION_TREE_HEIGHT: u32 = 5;
+const MAX_CONTRACT_FUNCTIONS: usize = 2_usize.pow(FUNCTION_TREE_HEIGHT);
+
 #[derive(Debug, Clone)]
 pub enum AztecMacroError {
     AztecNotFound,
     AztecComputeNoteHashAndNullifierNotFound { span: Span },
+    AztecContractHasTooManyFunctions { span: Span },
+    AztecContractConstructorMissing { span: Span },
 }
 
 impl From<AztecMacroError> for MacroError {
@@ -50,6 +55,16 @@ impl From<AztecMacroError> for MacroError {
                 secondary_message: None,
                 span: Some(span),
             },
+            AztecMacroError::AztecContractHasTooManyFunctions { span } => MacroError {
+                primary_message: format!("Contract can only have a maximum of {} functions", MAX_CONTRACT_FUNCTIONS),
+                secondary_message: None,
+                span: Some(span),
+            },
+            AztecMacroError::AztecContractConstructorMissing { span } => MacroError {
+                primary_message: "Contract must have a constructor function".to_owned(),
+                secondary_message: None,
+                span: Some(span),
+            },
         }
     }
 }
@@ -324,6 +339,23 @@ fn transform_module(
         }
     }
 
+    if module.functions.len() > MAX_CONTRACT_FUNCTIONS {
+        let crate_graph = &context.crate_graph[crate_id];
+        return Err((
+            AztecMacroError::AztecContractHasTooManyFunctions { span: Span::default() }.into(),
+            crate_graph.root_file_id,
+        ));
+    }
+
+    let constructor_defined = module.functions.iter().any(|func| func.name() == "constructor");
+    if !constructor_defined {
+        let crate_graph = &context.crate_graph[crate_id];
+        return Err((
+            AztecMacroError::AztecContractConstructorMissing { span: Span::default() }.into(),
+            crate_graph.root_file_id,
+        ));
+    }
+
     for func in module.functions.iter_mut() {
         for secondary_attribute in func.def.attributes.secondary.clone() {
             if is_custom_attribute(&secondary_attribute, "aztec(private)") {

From a632d74f99cd508b277b57e60bad880b9c6bcb42 Mon Sep 17 00:00:00 2001
From: Tom French <tom@tomfren.ch>
Date: Fri, 12 Jan 2024 17:04:23 +0000
Subject: [PATCH 2/2] fix: only run checks on aztec contracts

---
 aztec_macros/src/lib.rs | 39 ++++++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/aztec_macros/src/lib.rs b/aztec_macros/src/lib.rs
index 4bc96e3bfeb..a0ab4e79760 100644
--- a/aztec_macros/src/lib.rs
+++ b/aztec_macros/src/lib.rs
@@ -339,23 +339,6 @@ fn transform_module(
         }
     }
 
-    if module.functions.len() > MAX_CONTRACT_FUNCTIONS {
-        let crate_graph = &context.crate_graph[crate_id];
-        return Err((
-            AztecMacroError::AztecContractHasTooManyFunctions { span: Span::default() }.into(),
-            crate_graph.root_file_id,
-        ));
-    }
-
-    let constructor_defined = module.functions.iter().any(|func| func.name() == "constructor");
-    if !constructor_defined {
-        let crate_graph = &context.crate_graph[crate_id];
-        return Err((
-            AztecMacroError::AztecContractConstructorMissing { span: Span::default() }.into(),
-            crate_graph.root_file_id,
-        ));
-    }
-
     for func in module.functions.iter_mut() {
         for secondary_attribute in func.def.attributes.secondary.clone() {
             if is_custom_attribute(&secondary_attribute, "aztec(private)") {
@@ -372,6 +355,28 @@ fn transform_module(
             has_transformed_module = true;
         }
     }
+
+    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 crate_graph = &context.crate_graph[crate_id];
+            return Err((
+                AztecMacroError::AztecContractHasTooManyFunctions { span: Span::default() }.into(),
+                crate_graph.root_file_id,
+            ));
+        }
+
+        let constructor_defined = module.functions.iter().any(|func| func.name() == "constructor");
+        if !constructor_defined {
+            let crate_graph = &context.crate_graph[crate_id];
+            return Err((
+                AztecMacroError::AztecContractConstructorMissing { span: Span::default() }.into(),
+                crate_graph.root_file_id,
+            ));
+        }
+    }
+
     Ok(has_transformed_module)
 }