fix: free storage pointer used in param bug

huff_codegen/src/irgen/statements.rs

@@ -457,7 +457,7 @@ pub fn statement_gen(
                             )),
                             span: bf.span.clone(),
                             token: None,
-                        })
+                        });
                     }
 
                     let arg_index = bf.args[0].name.as_ref().unwrap();
@@ -476,7 +476,7 @@ pub fn statement_gen(
                             ),
                             span: bf.span.clone(),
                             token: None,
-                        })
+                        });
                     }
 
                     // Insert a 17 byte placeholder- will be filled when constructor args are added

huff_core/tests/free_storage_pointer.rs

@@ -0,0 +1,37 @@
+use huff_codegen::Codegen;
+use huff_lexer::Lexer;
+use huff_parser::Parser;
+use huff_utils::{prelude::FullFileSource, token::Token};
+
+/// Check that free storage pointers referenced outside of macro bodies
+/// are assigned correctly at compilation
+#[test]
+fn test_set_free_storage_pointers() {
+    let source: &str = r#"
+        #define constant FREE = FREE_STORAGE_POINTER()
+
+        #define macro TEST_MACRO(slot) = {
+            <slot> sload 0x00 mstore 0x20 0x00 return
+        }
+
+        #define macro MAIN() = {
+            TEST_MACRO(FREE)
+        }
+    "#;
+
+    // Parse tokens
+    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);
+
+    // Parse AST
+    let mut contract = parser.parse().unwrap();
+
+    // Derive storage pointers
+    contract.derive_storage_pointers();
+
+    // Assert the Free storage pointer has been set to 0
+    let mbytes = Codegen::generate_main_bytecode(&contract, None).unwrap();
+    assert!(mbytes.starts_with("6000"));
+}

huff_core/tests/macro_invoc_args.rs

@@ -48,10 +48,9 @@ fn test_all_opcodes_in_macro_args() {
             }}
 
             #define macro MAIN() = takes(0) returns(0) {{
-                RETURN1({})
+                RETURN1({o})
             }}
-        "#,
-            o
+        "#
         );
 
         // Lex + Parse

huff_core/tests/parser_errors.rs

@@ -169,9 +169,8 @@ fn test_invalid_token_in_macro_body() {
     for (value, kind) in invalids {
         let source = &format!(
             r#"#define macro CONSTANT() = takes (0) returns (0) {{
-            {}
-        }}"#,
-            value
+            {value}
+        }}"#
         );
 
         let full_source = FullFileSource { source, file: None, spans: vec![] };
@@ -215,9 +214,8 @@ fn test_invalid_token_in_label_definition() {
         let source = &format!(
             r#"#define macro CONSTANT() = takes (0) returns (0) {{
             lab:
-                {}
-        }}"#,
-            value
+                {value}
+        }}"#
         );
 
         let full_source = FullFileSource { source, file: None, spans: vec![] };
@@ -277,8 +275,7 @@ fn test_invalid_single_arg() {
                     e,
                     ParserError {
                         kind: ParserErrorKind::InvalidSingleArg(TokenKind::Ident(format!(
-                            "{}",
-                            random_char
+                            "{random_char}"
                         ))),
                         hint: Some("Expected number representing stack item count.".to_string()),
                         spans: AstSpan(vec![Span { start: 34, end: 35, file: None }]),

huff_lexer/tests/imports.rs

@@ -16,13 +16,12 @@ fn commented_lex_imports() {
     let import_str = "../huff-examples/erc20/contracts/utils/Ownable.huff";
     let source = format!(
         r#"
-    // #include "{}"
-    /* #include "{}" */
+    // #include "{import_str}"
+    /* #include "{import_str}" */
     /* test test test */
     #define macro ()
-    #include "{}"
-    "#,
-        import_str, import_str, import_str
+    #include "{import_str}"
+    "#
     );
 
     let lexed_imports = Lexer::lex_imports(&source);
@@ -35,13 +34,12 @@ fn multiple_lex_imports() {
     let import_str = "../huff-examples/erc20/contracts/utils/Ownable.huff";
     let source = format!(
         r#"
-    #include "{}"
-    #include "{}"
+    #include "{import_str}"
+    #include "{import_str}"
     /* test test test */
     #define macro ()
-    #include "{}"
-    "#,
-        import_str, import_str, import_str
+    #include "{import_str}"
+    "#
     );
 
     let lexed_imports = Lexer::lex_imports(&source);
@@ -56,10 +54,9 @@ fn multiple_lex_imports_single_quotes() {
     let import_str = "../huff-examples/erc20/contracts/utils/Ownable.huff";
     let source = format!(
         r#"
-    #include '{}'
-    #include '{}'
-    "#,
-        import_str, import_str
+    #include '{import_str}'
+    #include '{import_str}'
+    "#
     );
 
     let lexed_imports = Lexer::lex_imports(&source);

huff_parser/src/lib.rs

@@ -119,7 +119,7 @@ impl Parser {
                             kind: ParserErrorKind::InvalidDefinition(self.current_token.kind.clone()),
                             hint: Some("Definition must be one of: `function`, `event`, `constant`, `error`, `macro`, `fn`, or `test`.".to_string()),
                             spans: AstSpan(vec![self.current_token.span.clone()]),
-                        })
+                        });
                     }
                 };
             } else {
@@ -574,11 +574,10 @@ impl Parser {
                                     return Err(ParserError {
                                         kind: ParserErrorKind::InvalidPush(o),
                                         hint: Some(format!(
-                                            "Literal {:?} contains too many bytes for opcode \"{:?}\"",
-                                            hex_literal, o
+                                            "Literal {hex_literal:?} contains too many bytes for opcode \"{o:?}\""
                                         )),
                                         spans: AstSpan(curr_spans),
-                                    })
+                                    });
                                 }
 
                                 // Otherwise we can push the literal
@@ -924,8 +923,7 @@ impl Parser {
                                         self.current_token.kind.clone(),
                                     ),
                                     hint: Some(format!(
-                                        "Argument names cannot be EVM types: {}",
-                                        arg_str
+                                        "Argument names cannot be EVM types: {arg_str}"
                                     )),
                                     spans: AstSpan(vec![self.current_token.span.clone()]),
                                 })

huff_parser/tests/opcodes.rs

@@ -10,12 +10,11 @@ fn not_mistaken_as_opcode() {
             r#"
                 #define macro IS_AUTHORIZED(some_arg) = takes(0) returns(0) {{}}
                 #define macro MAIN() = takes(0) returns(0) {{
-                    IS_AUTHORIZED({})
-                    {}:
+                    IS_AUTHORIZED({label})
+                    {label}:
                         return
                 }}
-            "#,
-            label, label
+            "#
         );
         let flattened_source = FullFileSource { source, file: None, spans: vec![] };
         let lexer = Lexer::new(flattened_source);

huff_utils/src/ast.rs

@@ -6,7 +6,7 @@ use crate::{
     bytes_util::*,
     error::CodegenError,
     evm::Opcode,
-    prelude::{Span, TokenKind},
+    prelude::{MacroArg::Ident, Span, TokenKind},
 };
 use std::{
     collections::BTreeMap,
@@ -201,51 +201,47 @@ impl Contract {
         checking_constructor: bool,
     ) {
         let mut statements = macro_def.statements.clone();
+
         let mut i = 0;
         loop {
             if i >= statements.len() {
                 break
             }
             match &statements[i].clone().ty {
                 StatementType::Constant(const_name) => {
-                    tracing::debug!(target: "ast", "Found constant \"{}\" in macro def \"{}\" statements!", const_name, macro_def.name);
-                    if storage_pointers
-                        .iter()
-                        .filter(|pointer| pointer.0.eq(const_name))
-                        .collect::<Vec<&(String, [u8; 32])>>()
-                        .get(0)
-                        .is_none()
-                    {
-                        tracing::debug!(target: "ast", "No storage pointer already set for \"{}\"!", const_name);
-                        // Get the associated constant
-                        match self
-                            .constants
-                            .lock()
-                            .unwrap()
-                            .iter()
-                            .filter(|c| c.name.eq(const_name))
-                            .collect::<Vec<&ConstantDefinition>>()
-                            .get(0)
-                        {
-                            Some(c) => {
-                                let new_value = match c.value {
-                                    ConstVal::Literal(l) => l,
-                                    ConstVal::FreeStoragePointer(_) => {
-                                        let old_p = *last_p;
-                                        *last_p += 1;
-                                        str_to_bytes32(&format!("{old_p}"))
-                                    }
-                                };
-                                storage_pointers.push((const_name.to_string(), new_value));
-                            }
-                            None => {
-                                tracing::warn!(target: "ast", "CONSTANT \"{}\" NOT FOUND IN AST CONSTANTS", const_name)
-                            }
-                        }
-                    }
+                    self.assign_free_storage_pointers(
+                        const_name,
+                        &macro_def.name,
+                        storage_pointers,
+                        last_p,
+                    );
                 }
                 StatementType::MacroInvocation(mi) => {
                     tracing::debug!(target: "ast", "Found macro invocation: \"{}\" in macro def: \"{}\"!", mi.macro_name, macro_def.name);
+
+                    // Check for constant references in macro arguments
+                    let mut constant_args: Vec<String> = Vec::new();
+                    for arg in &mi.args {
+                        // check if it is a constant
+                        if let Ident(name) = arg {
+                            self.constants.lock().unwrap().iter().for_each(|constant| {
+                                if name == &constant.name {
+                                    tracing::debug!(target: "ast", "CONSTANT FOUND AS MACRO PARAMETER {}", name);
+                                    constant_args.push(name.to_string());
+                                }
+                            })
+                        }
+                    }
+                    // Assign constants that reference the Free Storage Pointer
+                    for constant_arg in constant_args {
+                        self.assign_free_storage_pointers(
+                            &constant_arg,
+                            &macro_def.name,
+                            storage_pointers,
+                            last_p,
+                        );
+                    }
+
                     match self
                         .macros
                         .iter()
@@ -321,6 +317,50 @@ impl Contract {
         // }
     }
 
+    fn assign_free_storage_pointers(
+        &self,
+        const_name: &String,
+        macro_name: &String,
+        storage_pointers: &mut Vec<(String, [u8; 32])>,
+        last_p: &mut i32,
+    ) {
+        tracing::debug!(target: "ast", "Found constant \"{}\" in macro def \"{}\" statements!", const_name, macro_name);
+        if storage_pointers
+            .iter()
+            .filter(|pointer| pointer.0.eq(const_name))
+            .collect::<Vec<&(String, [u8; 32])>>()
+            .get(0)
+            .is_none()
+        {
+            tracing::debug!(target: "ast", "No storage pointer already set for \"{}\"!", const_name);
+            // Get the associated constant
+            match self
+                .constants
+                .lock()
+                .unwrap()
+                .iter()
+                .filter(|c| c.name.eq(const_name))
+                .collect::<Vec<&ConstantDefinition>>()
+                .get(0)
+            {
+                Some(c) => {
+                    let new_value = match c.value {
+                        ConstVal::Literal(l) => l,
+                        ConstVal::FreeStoragePointer(_) => {
+                            let old_p = *last_p;
+                            *last_p += 1;
+                            str_to_bytes32(&format!("{old_p}"))
+                        }
+                    };
+                    storage_pointers.push((const_name.to_string(), new_value));
+                }
+                None => {
+                    tracing::warn!(target: "ast", "CONSTANT \"{}\" NOT FOUND IN AST CONSTANTS", const_name)
+                }
+            }
+        }
+    }
+
     /// Add override constants to the AST
     ///
     /// ## Overview