From 5579da90f4e35d6e41ad517ca58bfe5983b697e2 Mon Sep 17 00:00:00 2001
From: PraneshASP <praneshas2000@gmail.com>
Date: Thu, 24 Aug 2023 03:29:19 +0800
Subject: [PATCH] fix: restrict duplicate macros

---
 huff_parser/src/lib.rs     |  7 ++++
 huff_parser/tests/macro.rs | 66 ++++++++++++++++++++++++++++++++++++++
 huff_utils/src/error.rs    | 10 ++++++
 3 files changed, 83 insertions(+)

diff --git a/huff_parser/src/lib.rs b/huff_parser/src/lib.rs
index 48e88c0e..709bc691 100644
--- a/huff_parser/src/lib.rs
+++ b/huff_parser/src/lib.rs
@@ -104,6 +104,13 @@ impl Parser {
                     TokenKind::Macro | TokenKind::Fn | TokenKind::Test => {
                         let m = self.parse_macro()?;
                         tracing::info!(target: "parser", "SUCCESSFULLY PARSED MACRO {}", m.name);
+                        if contract.macros.iter().any(|existing| existing.name == m.name) {
+                            return Err(ParserError {
+                                kind: ParserErrorKind::DuplicateMacro(m.name),
+                                hint: Some("MACRO names should be unique".to_string()),
+                                spans: AstSpan(vec![self.spans[2].clone()]),
+                            })
+                        }
                         contract.macros.push(m);
                     }
                     TokenKind::JumpTable | TokenKind::JumpTablePacked | TokenKind::CodeTable => {
diff --git a/huff_parser/tests/macro.rs b/huff_parser/tests/macro.rs
index bc34e903..412d2136 100644
--- a/huff_parser/tests/macro.rs
+++ b/huff_parser/tests/macro.rs
@@ -1232,3 +1232,69 @@ fn empty_test_with_multi_flag_decorator() {
     assert_eq!(macro_definition, expected);
     assert_eq!(parser.current_token.kind, TokenKind::Eof);
 }
+
+#[test]
+fn test_duplicate_macro_error() {
+    let source = r#"
+    #define macro CONSTRUCTOR() = takes(0) returns (0) {}
+
+    #define macro MINT() = takes(0) returns (0) {
+        0x04 calldataload   // [to]
+        0x00                // [from (0x00), to]
+        0x24 calldataload   // [value, from, to]
+    }
+
+    #define macro MINT() = takes(0) returns (0) {
+        0x04 calldataload   // [to]
+        0x00                // [from (0x00), to]
+        0x24 calldataload   // [value, from, to]
+    }
+
+    #define macro MAIN() = takes(0) returns (0) {
+        0x00 calldataload 0xE0 shr
+        dup1 0x40c10f19 eq mints jumpi
+
+        mints:
+            MINT()
+    }
+    "#;
+
+    //let const_start = source.find("MINT()").unwrap_or(0);
+    //let const_end = const_start + "MINT()".len() - 1;
+    let mut start = 0;
+    let target = "MINT";
+    let target_len = target.len();
+    let mut occurrences = Vec::new();
+
+    while let Some(pos) = source[start..].find(target) {
+        let adjusted_start = start + pos;
+        let adjusted_end = adjusted_start + target_len - 1;
+
+        occurrences.push((adjusted_start, adjusted_end));
+        start = adjusted_end + 1;
+    }
+
+    let full_source = FullFileSource { source, file: None, spans: vec![] };
+    let lexer = Lexer::new(full_source.source);
+    let tokens = lexer.into_iter().map(|x| x.unwrap()).collect::<Vec<Token>>();
+    let mut parser = Parser::new(tokens, Some("".to_string()));
+
+    // This should be caught before codegen invalid macro statement
+    match parser.parse() {
+        Ok(_) => panic!("moose"),
+        Err(e) => {
+            assert_eq!(
+                e,
+                ParserError {
+                    kind: ParserErrorKind::DuplicateMacro("MINT".to_string()),
+                    hint: Some("MACRO names should be unique".to_string()),
+                    spans: AstSpan(vec![Span {
+                        start: occurrences[1].0.clone(),
+                        end: occurrences[1].1.clone(),
+                        file: None
+                    }]),
+                }
+            )
+        }
+    }
+}
diff --git a/huff_utils/src/error.rs b/huff_utils/src/error.rs
index 9ecb2631..b1d1af81 100644
--- a/huff_utils/src/error.rs
+++ b/huff_utils/src/error.rs
@@ -63,6 +63,8 @@ pub enum ParserErrorKind {
     InvalidDecoratorFlag(String),
     /// Invalid decorator flag argument
     InvalidDecoratorFlagArg(TokenKind),
+    /// Duplicate MACRO
+    DuplicateMacro(String),
 }
 
 /// A Lexing Error
@@ -488,6 +490,14 @@ impl fmt::Display for CompilerError {
                         pe.spans.error(pe.hint.as_ref())
                     )
                 }
+                ParserErrorKind::DuplicateMacro(mn) => {
+                    write!(
+                        f,
+                        "\nError: Duplicate MACRO name found: \"{}\" \n{}\n",
+                        mn,
+                        pe.spans.error(pe.hint.as_ref())
+                    )
+                }
             },
             CompilerError::PathBufRead(os_str) => {
                 write!(