From 30c2ca4a689ce55242a19c57d3918be388f799af Mon Sep 17 00:00:00 2001
From: Linda Guiga <lindaguiga3@gmail.com>
Date: Wed, 14 Aug 2024 10:51:34 +0100
Subject: [PATCH] Allow features with macros within macros

---
 evm_arithmetization/src/cpu/kernel/assembler.rs |  3 +++
 evm_arithmetization/src/cpu/kernel/parser.rs    | 14 +++++++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/evm_arithmetization/src/cpu/kernel/assembler.rs b/evm_arithmetization/src/cpu/kernel/assembler.rs
index c055e7a84..1bf97e604 100644
--- a/evm_arithmetization/src/cpu/kernel/assembler.rs
+++ b/evm_arithmetization/src/cpu/kernel/assembler.rs
@@ -219,6 +219,9 @@ fn expand_macros(
                     expanded.extend(expand_macros(body.clone(), macros, macro_counter));
                 }
             }
+            Item::ConditionalBlock(_, items) => {
+                expanded.extend(expand_macros(items.clone(), macros, macro_counter));
+            }
             item => {
                 expanded.push(item);
             }
diff --git a/evm_arithmetization/src/cpu/kernel/parser.rs b/evm_arithmetization/src/cpu/kernel/parser.rs
index 5e880964e..c661dabd2 100644
--- a/evm_arithmetization/src/cpu/kernel/parser.rs
+++ b/evm_arithmetization/src/cpu/kernel/parser.rs
@@ -246,6 +246,14 @@ mod tests {
     #[test]
     fn test_feature() {
         let code = r#"
+        %macro bar_foo
+            #[cfg(feature = feature_1)]
+            {
+                %bar
+                PUSH 3
+                ADD
+            }
+        %endmacro
         #[cfg(feature = feature_1)]
         {
             %macro bar
@@ -260,7 +268,7 @@ mod tests {
 
             #[cfg(feature = feature_1)]
             {
-                %bar
+                %bar_foo
                 PUSH 1
             }
             PUSH 3
@@ -291,6 +299,8 @@ mod tests {
         %macro bar
             PUSH 2
             MUL
+            PUSH 3
+            ADD
         %endmacro
 
         global foo_1:
@@ -353,6 +363,8 @@ mod tests {
         %macro bar
             PUSH 2
             MUL
+            PUSH 3
+            ADD
         %endmacro
 
         global foo_1: