feat: Nested macro calls (#44)

Author: Mouradif

CHANGELOG.md

@@ -3,6 +3,8 @@
 # Huff Neo Compiler changelog
 
 ## [Unreleased]
+- Support nesting of macro calls e.g. `MACRO1(MACRO2(0x1, 0x2), 0x3)`. (See: #40)
+  - Thank you very much to @Mouradif for the contribution!
 
 ## [1.1.4] - 2025-03-17
 - Update dependencies to the latest version.

Cargo.toml

@@ -69,7 +69,7 @@ getrandom = { version = "0.3.1", features = ["wasm_js"] }
 rayon = "1.10.0"
 
 [profile.test]
-debug = 1
+debug = "full"
 incremental = true
 
 [profile.release]

book/huff-language/macros-and-functions.md

@@ -39,7 +39,7 @@ by the constructor.
 
 ### Macro Arguments
 
-Macros can accept arguments to be "called" inside the macro or passed as a reference. Macro arguments may be one of: label, opcode, literal, or a constant. Since macros are inlined at compile-time, the arguments are not evaluated at runtime and are instead inlined as well.
+Macros can accept arguments, which can be used within the macro itself or passed as reference. These arguments can be labels, opcodes, literals, constants, or other macro calls. Since macros are inlined at compile time, their arguments are also inlined and not evaluated at runtime.
 
 #### Example
 
@@ -137,7 +137,7 @@ your contract, and it is essentially a trade-off of decreasing contract size
 for a small extra runtime gas cost (`22 + n_inputs * 3 + n_outputs * 3` gas
 per invocation, to be exact).
 
-Functions are one of the only high-level abstractions
+Functions are one of the few high-level abstractions
 in Huff, so it is important to understand what the compiler adds to your code
 when they are utilized. It is not always beneficial to re-use code, especially
 if it is a small / inexpensive set of operations. However, for larger contracts

crates/codegen/src/irgen/arg_calls.rs

@@ -2,7 +2,7 @@ use crate::Codegen;
 use huff_neo_utils::ast::span::AstSpan;
 use huff_neo_utils::prelude::*;
 use std::str::FromStr;
-// Arguments can be literals, labels, opcodes, or constants
+// Arguments can be literals, labels, opcodes, constants, or macro calls
 // !! IF THERE IS AMBIGUOUS NOMENCLATURE
 // !! (E.G. BOTH OPCODE AND LABEL ARE THE SAME STRING)
 // !! COMPILATION _WILL_ ERROR
@@ -136,6 +136,40 @@ pub fn bubble_arg_call(
                             *offset += 3;
                         }
                     }
+                    MacroArg::MacroCall(inner_mi) => {
+                        tracing::debug!(target: "codegen", "Found MacroArg::MacroCall IN \"{}\" Macro Invocation: \"{}\"!", macro_invoc.1.macro_name, inner_mi.macro_name);
+
+                        if let Some(called_macro) = contract.find_macro_by_name(&inner_mi.macro_name) {
+                            tracing::debug!(target: "codegen", "Found valid macro: {}", called_macro.name);
+                            let mut new_scope = scope.to_vec();
+                            new_scope.push(called_macro);
+                            let mut new_mis = mis.to_vec();
+                            new_mis.push((starting_offset, inner_mi.clone()));
+                            match Codegen::macro_to_bytecode(
+                                evm_version,
+                                called_macro,
+                                contract,
+                                &mut new_scope,
+                                starting_offset,
+                                &mut new_mis,
+                                false,
+                                None,
+                            ) {
+                                Ok(expanded_macro) => {
+                                    bytes.extend(expanded_macro.bytes);
+                                }
+                                Err(e) => {
+                                    return Err(e);
+                                }
+                            }
+                        } else {
+                            return Err(CodegenError {
+                                kind: CodegenErrorKind::MissingMacroDefinition(inner_mi.macro_name.clone()),
+                                span: inner_mi.span.clone(),
+                                token: None,
+                            });
+                        }
+                    }
                 }
             } else {
                 tracing::warn!(target: "codegen", "\"{}\" FOUND IN MACRO DEF BUT NOT IN MACRO INVOCATION!", arg_name);

crates/core/tests/macro_invoc_args.rs

@@ -309,3 +309,155 @@ fn test_bubbled_arg_with_different_name() {
     // Check the bytecode
     assert_eq!(main_bytecode, expected_bytecode);
 }
+
+#[test]
+fn test_nested_macro_call_no_args() {
+    let source = r#"
+          #define macro COMPUTE_FIXED(a) = takes(0) returns(0) {
+            0x01
+            0x01
+            <a>
+          }
+
+          #define macro ADD() = takes(0) returns(0) {
+            add
+          }
+
+          #define macro SUB() = takes(0) returns(0) {
+            sub
+          }
+
+          #define macro MAIN() = takes(0) returns(0) {
+            COMPUTE_FIXED(ADD())
+            COMPUTE_FIXED(SUB())
+          }
+        "#;
+
+    // Lex + Parse
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+    let mut contract = parser.parse().unwrap();
+    contract.derive_storage_pointers();
+
+    let evm_version = EVMVersion::default();
+
+    // Create main and constructor bytecode
+    let main_bytecode = Codegen::generate_main_bytecode(&evm_version, &contract, None).unwrap();
+
+    // Full expected bytecode output (generated from huff-neo) (placed here as a reference)
+    let expected_bytecode = "60016001016001600103";
+
+    // Check the bytecode
+    assert_eq!(main_bytecode.to_lowercase(), expected_bytecode.to_lowercase());
+}
+
+#[test]
+fn test_nested_macro_calls() {
+    let source = r#"
+            #define macro ADD(a, b) = takes(0) returns(1) {
+                <a> <b> add
+            }
+
+            #define macro SUB(a, b) = takes(0) returns(1) {
+                <a> <b> sub
+            }
+
+            #define macro COMPUTE_AND_STORE(operation, offset) = takes(0) returns(0) {
+                <operation>
+                <offset> mstore
+            }
+
+            #define macro MAIN() = takes(0) returns(0) {
+                COMPUTE_AND_STORE(ADD(0x05, 0x06), 0x00)
+                COMPUTE_AND_STORE(SUB(0x05, 0x06), 0x20)
+            }
+        "#;
+
+    // Lex + Parse
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+    let mut contract = parser.parse().unwrap();
+    contract.derive_storage_pointers();
+
+    let evm_version = EVMVersion::default();
+
+    // Create main and constructor bytecode
+    let main_bytecode = Codegen::generate_main_bytecode(&evm_version, &contract, None).unwrap();
+
+    // Full expected bytecode output (generated from huff-neo) (placed here as a reference)
+    let expected_bytecode = "60056006016000526005600603602052";
+
+    // Check the bytecode
+    assert_eq!(main_bytecode.to_lowercase(), expected_bytecode.to_lowercase());
+}
+
+#[test]
+fn test_very_nested_macro_calls() {
+    let source = r#"
+            #define macro ADD(a, b) = takes(0) returns(1) {
+                <a> <b> add
+            }
+
+            #define macro SUB(a, b) = takes(0) returns(1) {
+                <a> <b> sub
+            }
+
+            #define macro COMPUTE_AND_STORE(operation, offset) = takes(0) returns(0) {
+                <operation> <offset> mstore
+            }
+
+            #define macro MAIN() = takes(0) returns(0) {
+                COMPUTE_AND_STORE(ADD(0x05, SUB(0x05, ADD(ADD(0x0a, 0x16), 0x04))), 0x00)
+            }
+        "#;
+
+    // Lex + Parse
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+    let mut contract = parser.parse().unwrap();
+    contract.derive_storage_pointers();
+
+    let evm_version = EVMVersion::default();
+
+    // Create main and constructor bytecode
+    let main_bytecode = Codegen::generate_main_bytecode(&evm_version, &contract, None).unwrap();
+
+    // Full expected bytecode output (generated from huff-neo) (placed here as a reference)
+    let expected_bytecode = "60056005600a6016016004010301600052";
+
+    // Check the bytecode
+    assert_eq!(main_bytecode.to_lowercase(), expected_bytecode.to_lowercase());
+}
+
+#[test]
+fn test_nested_macro_call_missing_macro_definition() {
+    let source = r#"
+          #define macro MACRO1(a) = takes(0) returns(0) {
+            <a>
+          }
+
+          #define macro MAIN() = takes(0) returns(0) {
+            MACRO1(MISSING_MACRO())
+          }
+        "#;
+
+    // Lex + Parse
+    let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(flattened_source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, None);
+    let mut contract = parser.parse().unwrap();
+    contract.derive_storage_pointers();
+
+    // Codegen should fail with an error
+    let codegen_result = Codegen::generate_main_bytecode(&EVMVersion::default(), &contract, None);
+
+    assert!(codegen_result.is_err());
+    assert_eq!(codegen_result.unwrap_err().kind, CodegenErrorKind::MissingMacroDefinition(String::from("MISSING_MACRO")));
+}

crates/parser/src/lib.rs

@@ -1033,8 +1033,25 @@ impl Parser {
                     self.consume();
                 }
                 TokenKind::Ident(ident) => {
-                    args.push(MacroArg::Ident(ident));
-                    self.consume();
+                    if self.peek().unwrap().kind == TokenKind::OpenParen {
+                        // It's a nested macro call
+                        let mut curr_spans = vec![self.current_token.span.clone()];
+                        self.match_kind(TokenKind::Ident("MACRO_NAME".to_string()))?;
+                        // Parse Macro Call
+                        let lit_args = self.parse_macro_call_args()?;
+                        // Grab all spans following our macro invocation spam
+                        if let Some(i) = self.spans.iter().position(|s| s.eq(&curr_spans[0])) {
+                            curr_spans.append(&mut self.spans[(i + 1)..].to_vec());
+                        }
+                        args.push(MacroArg::MacroCall(MacroInvocation {
+                            macro_name: ident,
+                            args: lit_args,
+                            span: AstSpan(curr_spans.clone()),
+                        }));
+                    } else {
+                        args.push(MacroArg::Ident(ident));
+                        self.consume();
+                    }
                 }
                 TokenKind::Calldata => {
                     args.push(MacroArg::Ident("calldata".to_string()));

crates/utils/src/ast/huff.rs

@@ -472,6 +472,8 @@ pub enum MacroArg {
     Ident(String),
     /// An Arg Call
     ArgCall(String),
+    /// A Nested Macro Call
+    MacroCall(MacroInvocation),
 }
 
 /// Free Storage Pointer Unit Struct