Repository: huff-language/huff-rs
Pull request #130 (2022-06-24 18:58:10)

fix(huff_lexer): label mistaken as opcode if prefixed with it
## Overview

If a label is prefixed with an opcode (example : `not_authed`), it will be lexed as an opcode followed by a colon.





Patch:
diff --git a/huff_codegen/src/irgen/statements.rs b/huff_codegen/src/irgen/statements.rs
index 4c6e432..1d61512 100644
--- a/huff_codegen/src/irgen/statements.rs
+++ b/huff_codegen/src/irgen/statements.rs
@@ -18,7 +18,7 @@ pub fn statement_gen(
 ) -> Result<Vec<(usize, Bytes)>, CodegenError> {
     let mut bytes = vec![];
 
-    tracing::debug!(target: "codegen", "Got Statement: {}", s.ty);
+    tracing::debug!(target: "codegen", "Got Statement: {:?}", s);
 
     match &s.ty {
         StatementType::MacroInvocation(mi) => {
@@ -58,7 +58,7 @@ pub fn statement_gen(
                 };
 
             // Set jump table values
-            tracing::debug!(target: "codegen", "Unmatched jumps: {:?}", res.unmatched_jumps.iter().map(|uj| uj.label.clone()).collect::<Vec<String>>());
+            tracing::debug!(target: "codegen", "Unmatched jumps: {:?}", res.unmatched_jumps);
             for j in res.unmatched_jumps.iter_mut() {
                 let new_index = j.bytecode_index;
                 j.bytecode_index = 0;
@@ -80,7 +80,7 @@ pub fn statement_gen(
         }
         StatementType::Label(label) => {
             // Add JUMPDEST opcode to final result and add to label_indices
-            tracing::info!(target: "codegen", "RECURSE BYTECODE GOT LABEL: {:?}", label.name);
+            tracing::info!(target: "codegen", "RECURSE BYTECODE GOT LABEL: {:?}", label);
             label_indices.insert(label.name.clone(), *offset);
             bytes.push((*offset, Bytes(Opcode::Jumpdest.to_string())));
             *offset += 1;
diff --git a/huff_core/src/lib.rs b/huff_core/src/lib.rs
index 85274e8..6ac535f 100644
--- a/huff_core/src/lib.rs
+++ b/huff_core/src/lib.rs
@@ -185,7 +185,7 @@ impl<'a> Compiler {
             file: Some(Arc::clone(&file)),
             spans: flattened.1,
         };
-        tracing::debug!(target: "core", "GOT FULL SOURCE FOR PATH: {:?}", file.path);
+        tracing::debug!(target: "core", "GOT FULL SOURCE: \"{:?}\"", full_source);
 
         // Perform Lexical Analysis
         // Create a new lexer from the FileSource, flattening dependencies
diff --git a/huff_lexer/src/lib.rs b/huff_lexer/src/lib.rs
index 79aef2f..f1c3a65 100644
--- a/huff_lexer/src/lib.rs
+++ b/huff_lexer/src/lib.rs
@@ -24,6 +24,8 @@ pub enum Context {
     MacroDefinition,
     /// Macro's body context
     MacroBody,
+    /// Macro's argument context (definition or being called)
+    MacroArgs,
     /// ABI context
     Abi,
     /// Lexing args of functions inputs/outputs and events
@@ -455,7 +457,7 @@ impl<'a> Iterator for Lexer<'a> {
 
                     // goes over all opcodes
                     for opcode in OPCODES {
-                        if self.context != Context::MacroBody {
+                        if self.context != Context::MacroBody || found_kind != None {
                             break
                         }
                         if opcode == pot_op {
@@ -564,14 +566,18 @@ impl<'a> Iterator for Lexer<'a> {
                 }
                 '=' => TokenKind::Assign,
                 '(' => {
-                    if self.context == Context::Abi {
-                        self.context = Context::AbiArgs;
+                    match self.context {
+                        Context::Abi => self.context = Context::AbiArgs,
+                        Context::MacroBody => self.context = Context::MacroArgs,
+                        _ => {}
                     }
                     TokenKind::OpenParen
                 }
                 ')' => {
-                    if self.context == Context::AbiArgs {
-                        self.context = Context::Abi;
+                    match self.context {
+                        Context::AbiArgs => self.context = Context::Abi,
+                        Context::MacroArgs => self.context = Context::MacroBody,
+                        _ => {}
                     }
                     TokenKind::CloseParen
                 }
diff --git a/huff_utils/src/ast.rs b/huff_utils/src/ast.rs
index de818e5..bddf519 100644
--- a/huff_utils/src/ast.rs
+++ b/huff_utils/src/ast.rs
@@ -8,11 +8,7 @@ use crate::{
     evm::Opcode,
     prelude::{Span, TokenKind},
 };
-use std::{
-    collections::BTreeMap,
-    fmt::{Display, Formatter},
-    path::PathBuf,
-};
+use std::{collections::BTreeMap, path::PathBuf};
 
 /// A contained literal
 pub type Literal = [u8; 32];
@@ -630,22 +626,3 @@ pub enum StatementType {
     /// A built-in function call
     BuiltinFunctionCall(BuiltinFunctionCall),
 }
-
-impl Display for StatementType {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        match self {
-            StatementType::Literal(l) => write!(f, "LITERAL: {}", bytes32_to_string(l, true)),
-            StatementType::Opcode(o) => write!(f, "OPCODE: {}", o),
-            StatementType::MacroInvocation(m) => {
-                write!(f, "MACRO INVOCATION: {}", m.macro_name)
-            }
-            StatementType::Constant(c) => write!(f, "CONSTANT: {}", c),
-            StatementType::ArgCall(c) => write!(f, "ARG CALL: {}", c),
-            StatementType::Label(l) => write!(f, "LABEL: {}", l.name),
-            StatementType::LabelCall(l) => write!(f, "LABEL CALL: {}", l),
-            StatementType::BuiltinFunctionCall(b) => {
-                write!(f, "BUILTIN FUNCTION CALL: {:?}", b.kind)
-            }
-        }
-    }
-}


Test patch:
diff --git a/huff_parser/tests/opcodes.rs b/huff_parser/tests/opcodes.rs
new file mode 100644
index 0000000..eed0d88
--- /dev/null
+++ b/huff_parser/tests/opcodes.rs
@@ -0,0 +1,36 @@
+use huff_lexer::*;
+use huff_parser::*;
+use huff_utils::{evm::Opcode, prelude::*};
+
+#[test]
+fn not_mistaken_as_opcode() {
+    for opcode in OPCODES {
+        let label = format!("{}_label", (*opcode).to_owned());
+        let source = &format!(
+            r#"
+                #define macro IS_AUTHORIZED(some_arg) = takes(0) returns(0) {{}}
+                #define macro MAIN() = takes(0) returns(0) {{
+                    IS_AUTHORIZED({})
+                    {}:
+                        return
+                }}
+            "#,
+            label, label
+        );
+        let flattened_source = FullFileSource { source, file: None, spans: vec![] };
+        let lexer = Lexer::new(flattened_source);
+
+        let tokens = lexer
+            .into_iter()
+            .map(|x| x.unwrap())
+            .filter(|x| !matches!(x.kind, TokenKind::Whitespace))
+            .collect::<Vec<Token>>();
+        let actual_label_arg = tokens[tokens.len() - 7].kind.clone();
+        let actual_label = tokens[tokens.len() - 5].kind.clone();
+        let mut parser = Parser::new(tokens, None);
+        // parsing to ensure tokens syntax is valid
+        let contract = parser.parse().unwrap();
+        assert_eq!(actual_label_arg, TokenKind::Ident(label.clone()));
+        assert_eq!(actual_label, TokenKind::Label(label));
+    }
+}