Add error for duplicate macro labels (Original author is 0xrusowsky) (#7)

cakevm · web-flow · commit 2e3d1a7915ff · 2025-01-08T21:18:18.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/codegen/README.md b/crates/codegen/README.md
@@ -94,6 +94,7 @@ let contract = Contract {
   functions: vec![],
   events: vec![],
   tables: vec![],
+  labels: HashSet::new(),
 };
 
 // Generate the main bytecode
@@ -149,6 +150,7 @@ let contract = Contract {
   functions: vec![],
   events: vec![],
   tables: vec![],
+  labels: HashSet::new(),
 };
 
 // Generate the constructor bytecode
diff --git a/crates/codegen/tests/abigen.rs b/crates/codegen/tests/abigen.rs
@@ -1,11 +1,11 @@
+use huff_neo_codegen::Codegen;
+use huff_neo_utils::prelude::*;
+use std::collections::HashSet;
 use std::{
     collections::BTreeMap,
     sync::{Arc, Mutex},
 };
 
-use huff_neo_codegen::Codegen;
-use huff_neo_utils::prelude::*;
-
 #[test]
 fn constructs_valid_abi() {
     let constructor = MacroDefinition {
@@ -28,6 +28,7 @@ fn constructs_valid_abi() {
         functions: vec![],
         events: vec![],
         tables: vec![],
+        labels: HashSet::new(),
     };
 
     // Generate the abi from the contract
@@ -68,6 +69,7 @@ fn missing_constructor_fails() {
         functions: vec![],
         events: vec![],
         tables: vec![],
+        labels: HashSet::new(),
     };
 
     // Generate the abi from the contract
diff --git a/crates/lexer/Cargo.toml b/crates/lexer/Cargo.toml
@@ -12,5 +12,7 @@ keywords = ["bytecode", "compiler", "evm", "huff", "rust"]
 
 [dependencies]
 huff-neo-utils.workspace = true
+
+lazy_static.workspace = true
 regex.workspace = true
 tracing.workspace = true
diff --git a/crates/lexer/src/lib.rs b/crates/lexer/src/lib.rs
@@ -1,11 +1,37 @@
 use huff_neo_utils::prelude::*;
+use lazy_static::lazy_static;
 use regex::Regex;
+use std::collections::HashMap;
 use std::{
     iter::{Peekable, Zip},
     ops::RangeFrom,
     str::Chars,
 };
 
+lazy_static! {
+    static ref TOKEN: HashMap<String, TokenKind> = HashMap::from_iter(vec![
+        (TokenKind::Macro.to_string(), TokenKind::Macro),
+        (TokenKind::Fn.to_string(), TokenKind::Fn),
+        (TokenKind::Test.to_string(), TokenKind::Test),
+        (TokenKind::Function.to_string(), TokenKind::Function),
+        (TokenKind::Constant.to_string(), TokenKind::Constant),
+        (TokenKind::Error.to_string(), TokenKind::Error),
+        (TokenKind::Takes.to_string(), TokenKind::Takes),
+        (TokenKind::Returns.to_string(), TokenKind::Returns),
+        (TokenKind::Event.to_string(), TokenKind::Event),
+        (TokenKind::NonPayable.to_string(), TokenKind::NonPayable),
+        (TokenKind::Payable.to_string(), TokenKind::Payable),
+        (TokenKind::Indexed.to_string(), TokenKind::Indexed),
+        (TokenKind::View.to_string(), TokenKind::View),
+        (TokenKind::Pure.to_string(), TokenKind::Pure),
+        // First check for packed jump table
+        (TokenKind::JumpTablePacked.to_string(), TokenKind::JumpTablePacked),
+        // Match with jump table if not
+        (TokenKind::JumpTable.to_string(), TokenKind::JumpTable),
+        (TokenKind::CodeTable.to_string(), TokenKind::CodeTable),
+    ]);
+}
+
 /// Defines a context in which the lexing happens.
 /// Allows to differentiate between EVM types and opcodes that can either
 /// be identical or the latter being a substring of the former (example : bytes32 and byte)
@@ -168,37 +194,9 @@ impl<'a> Lexer<'a> {
                     let (word, start, mut end) = self.eat_while(Some(ch), |c| c.is_alphanumeric() || c == '_');
 
                     let mut found_kind: Option<TokenKind> = None;
-                    let keys = [
-                        TokenKind::Macro,
-                        TokenKind::Fn,
-                        TokenKind::Test,
-                        TokenKind::Function,
-                        TokenKind::Constant,
-                        TokenKind::Error,
-                        TokenKind::Takes,
-                        TokenKind::Returns,
-                        TokenKind::Event,
-                        TokenKind::NonPayable,
-                        TokenKind::Payable,
-                        TokenKind::Indexed,
-                        TokenKind::View,
-                        TokenKind::Pure,
-                        // First check for packed jump table
-                        TokenKind::JumpTablePacked,
-                        // Match with jump table if not
-                        TokenKind::JumpTable,
-                        TokenKind::CodeTable,
-                    ];
-                    for kind in keys.into_iter() {
-                        if self.context == Context::MacroBody {
-                            break;
-                        }
-                        let key = kind.to_string();
-                        let peeked = word.clone();
-
-                        if key == peeked {
-                            found_kind = Some(kind);
-                            break;
+                    if self.context != Context::MacroBody {
+                        if let Some(kind) = TOKEN.get(&word) {
+                            found_kind = Some(kind.clone());
                         }
                     }
 
@@ -209,6 +207,7 @@ impl<'a> Lexer<'a> {
                         found_kind = None;
                     }
 
+                    // Set the context based on the found token kind
                     if let Some(kind) = &found_kind {
                         match kind {
                             TokenKind::Macro | TokenKind::Fn | TokenKind::Test => self.context = Context::MacroDefinition,
@@ -437,10 +436,9 @@ impl<'a> Lexer<'a> {
         let (integer_str, start, end) = self.eat_while(Some(initial_char), |ch| ch.is_ascii_digit());
 
         let integer = integer_str.parse().unwrap();
-
         let integer_token = TokenKind::Num(integer);
-        let span = self.source.relative_span_by_pos(start, end);
-        Ok(Token { kind: integer_token, span })
+
+        Ok(Token { kind: integer_token, span: self.source.relative_span_by_pos(start, end) })
     }
 
     fn eat_hex_digit(&mut self, initial_char: char) -> TokenResult {
@@ -461,8 +459,8 @@ impl<'a> Lexer<'a> {
         };
 
         start += 2;
-        let span = self.source.relative_span_by_pos(start, end);
-        Ok(Token { kind, span })
+
+        Ok(Token { kind, span: self.source.relative_span_by_pos(start, end) })
     }
 
     /// Skips white space. They are not significant in the source language
diff --git a/crates/parser/README.md b/crates/parser/README.md
@@ -58,6 +58,7 @@ let expected_contract = Contract {
   functions: vec![],
   events: vec![],
   tables: vec![],
+  labels: HashSet::new(),
 };
 assert_eq!(unwrapped_contract.macros, expected_contract.macros);
 ```
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
@@ -68,7 +68,7 @@ impl Parser {
             }
             // Check for a decorator above a test macro
             else if self.check(TokenKind::Pound) {
-                let m = self.parse_macro()?;
+                let m = self.parse_macro(&mut contract)?;
                 tracing::info!(target: "parser", "SUCCESSFULLY PARSED MACRO {}", m.name);
                 contract.macros.push(m);
             }
@@ -100,7 +100,7 @@ impl Parser {
                         contract.errors.push(e);
                     }
                     TokenKind::Macro | TokenKind::Fn | TokenKind::Test => {
-                        let m = self.parse_macro()?;
+                        let m = self.parse_macro(&mut contract)?;
                         tracing::info!(target: "parser", "SUCCESSFULLY PARSED MACRO {}", m.name);
                         self.check_duplicate_macro(&contract, &m)?;
                         contract.macros.push(m);
@@ -187,6 +187,24 @@ impl Parser {
         std::mem::discriminant(&self.current_token.kind) == std::mem::discriminant(&kind)
     }
 
+    /// Checks whether the input label is unique in a macro.
+    fn check_duplicate_label(&self, contract: &mut Contract, macro_name: &str, label: String, span: AstSpan) -> Result<(), ParserError> {
+        let key = format!("{macro_name}{label}");
+        if contract.labels.contains(&key) {
+            println!("DUPLICATED LABEL NAME: {:?}", span);
+            tracing::error!(target: "parser", "DUPLICATED LABEL NAME: {}", label);
+            Err(ParserError {
+                kind: ParserErrorKind::DuplicateLabel(label.clone()),
+                hint: Some(format!("Duplicated label name: \"{label}\" in macro: \"{macro_name}\"")),
+                spans: span,
+                cursor: self.cursor,
+            })
+        } else {
+            contract.labels.insert(key);
+            Ok(())
+        }
+    }
+
     /// Checks if there is a duplicate macro name
     pub fn check_duplicate_macro(&self, contract: &Contract, m: &MacroDefinition) -> Result<(), ParserError> {
         if contract.macros.binary_search_by(|_macro| _macro.name.cmp(&m.name)).is_ok() {
@@ -489,7 +507,7 @@ impl Parser {
     /// Parses a macro.
     ///
     /// It should parse the following : macro MACRO_NAME(args...) = takes (x) returns (n) {...}
-    pub fn parse_macro(&mut self) -> Result<MacroDefinition, ParserError> {
+    pub fn parse_macro(&mut self, contract: &mut Contract) -> Result<MacroDefinition, ParserError> {
         let mut decorator: Option<Decorator> = None;
         if self.check(TokenKind::Pound) {
             decorator = Some(self.parse_decorator()?);
@@ -518,7 +536,7 @@ impl Parser {
         let macro_takes = self.match_kind(TokenKind::Takes).map_or(Ok(0), |_| self.parse_single_arg())?;
         let macro_returns = self.match_kind(TokenKind::Returns).map_or(Ok(0), |_| self.parse_single_arg())?;
 
-        let macro_statements: Vec<Statement> = self.parse_body()?;
+        let macro_statements: Vec<Statement> = self.parse_body(&macro_name, contract)?;
 
         Ok(MacroDefinition::new(
             macro_name,
@@ -536,7 +554,7 @@ impl Parser {
     /// Parse the body of a macro.
     ///
     /// Only HEX, OPCODES, labels, builtins, and MACRO calls should be authorized.
-    pub fn parse_body(&mut self) -> Result<Vec<Statement>, ParserError> {
+    pub fn parse_body(&mut self, macro_name: &str, contract: &mut Contract) -> Result<Vec<Statement>, ParserError> {
         let mut statements: Vec<Statement> = Vec::new();
         self.match_kind(TokenKind::OpenBrace)?;
         tracing::info!(target: "parser", "PARSING MACRO BODY");
@@ -618,6 +636,7 @@ impl Parser {
                 TokenKind::Label(l) => {
                     let mut curr_spans = vec![self.current_token.span.clone()];
                     self.consume();
+                    self.check_duplicate_label(contract, macro_name, l.to_string(), AstSpan(curr_spans.clone()))?;
                     let inner_statements: Vec<Statement> = self.parse_label()?;
                     inner_statements.iter().for_each(|a| curr_spans.extend_from_slice(a.span.inner_ref()));
                     tracing::info!(target: "parser", "PARSED LABEL \"{}\" INSIDE MACRO WITH {} STATEMENTS.", l, inner_statements.len());
diff --git a/crates/parser/tests/labels.rs b/crates/parser/tests/labels.rs
@@ -1,6 +1,6 @@
 use huff_neo_lexer::*;
 use huff_neo_parser::*;
-use huff_neo_utils::{evm::Opcode, prelude::*};
+use huff_neo_utils::{error::ParserErrorKind, evm::Opcode, prelude::*};
 
 #[test]
 fn multiline_labels() {
@@ -382,3 +382,25 @@ pub fn builtins_under_labels() {
         assert_eq!(s.span, md_expected.statements[i].span);
     }
 }
+
+#[test]
+fn duplicated_labels() {
+    let source = r#"
+    #define macro MAIN() = takes(0) returns(0) {
+        cool_label jump
+        cool_label jump
+        cool_label: 0x00
+        dup_label: 0x00
+        dup_label: 0x00
+    }
+    "#;
+    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);
+
+    // Grab the first macro
+    let parse_result = parser.parse();
+    assert!(parse_result.is_err());
+    assert_eq!(parse_result.unwrap_err().kind, ParserErrorKind::DuplicateLabel("dup_label".to_string()));
+}
diff --git a/crates/utils/src/ast.rs b/crates/utils/src/ast.rs
@@ -10,6 +10,7 @@ use crate::{
     evm_version::EVMVersion,
     prelude::{MacroArg::Ident, Span, TokenKind},
 };
+use std::collections::HashSet;
 use std::{
     collections::BTreeMap,
     fmt::{Display, Formatter},
@@ -112,6 +113,8 @@ pub struct Contract {
     pub events: Vec<EventDefinition>,
     /// Tables
     pub tables: Vec<TableDefinition>,
+    /// Labels
+    pub labels: HashSet<String>,
 }
 
 impl Contract {
diff --git a/crates/utils/src/error.rs b/crates/utils/src/error.rs
@@ -65,6 +65,8 @@ pub enum ParserErrorKind {
     InvalidDecoratorFlag(String),
     /// Invalid decorator flag argument
     InvalidDecoratorFlagArg(TokenKind),
+    /// Duplicate label
+    DuplicateLabel(String),
     /// Duplicate MACRO
     DuplicateMacro(String),
 }
@@ -413,6 +415,9 @@ impl fmt::Display for CompilerError {
                         pe.spans.error(pe.hint.as_ref())
                     )
                 }
+                ParserErrorKind::DuplicateLabel(label) => {
+                    write!(f, "\nError: Duplicate label: \"{}\" \n{}\n", label, 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()))
                 }
diff --git a/crates/utils/src/files.rs b/crates/utils/src/files.rs
@@ -343,6 +343,12 @@ impl Span {
                 f.source
                     .as_ref()
                     .map(|s| {
+                        if self.start >= s.len() {
+                            return "\nInternal compiler error: Start index out of range".to_string();
+                        }
+                        if self.end >= s.len() {
+                            return "\nInternal compiler error: End index out of range: file".to_string();
+                        }
                         let line_num = &s[0..self.start].as_bytes().iter().filter(|&&c| c == b'\n').count() + 1;
                         let line_start = &s[0..self.start].rfind('\n').unwrap_or(0);
                         let line_end = self.end + s[self.end..s.len()].find('\n').unwrap_or(s.len() - self.end).to_owned();

Original file line number	Diff line number	Diff line change
`@@ -65,6 +65,8 @@ pub enum ParserErrorKind {`
`65`	`65`	`InvalidDecoratorFlag(String),`
`66`	`66`	`/// Invalid decorator flag argument`
`67`	`67`	`InvalidDecoratorFlagArg(TokenKind),`
	`68`	`+ /// Duplicate label`
	`69`	`+ DuplicateLabel(String),`
`68`	`70`	`/// Duplicate MACRO`
`69`	`71`	`DuplicateMacro(String),`
`70`	`72`	`}`
`@@ -413,6 +415,9 @@ impl fmt::Display for CompilerError {`
`413`	`415`	`pe.spans.error(pe.hint.as_ref())`
`414`	`416`	`)`
`415`	`417`	`}`
	`418`	`+ ParserErrorKind::DuplicateLabel(label) => {`
	`419`	`+ write!(f, "\nError: Duplicate label: \"{}\" \n{}\n", label, pe.spans.error(pe.hint.as_ref()))`
	`420`	`+ }`
`416`	`421`	`ParserErrorKind::DuplicateMacro(mn) => {`
`417`	`422`	`write!(f, "\nError: Duplicate MACRO name found: \"{}\" \n{}\n", mn, pe.spans.error(pe.hint.as_ref()))`
`418`	`423`	`}`