Repository: huff-language/huff-rs
Pull request #181 (2022-07-29 01:31:46)

Fix builtins under jumplabels
## Overview

Fixes the parser not recognizing builtin function calls underneath jumplabels.


Patch:
diff --git a/huff_parser/src/lib.rs b/huff_parser/src/lib.rs
index f559334..b617489 100644
--- a/huff_parser/src/lib.rs
+++ b/huff_parser/src/lib.rs
@@ -450,7 +450,7 @@ impl Parser {
 
     /// Parse the body of a macro.
     ///
-    /// Only HEX, OPCODES, labels and MACRO calls should be authorized.
+    /// Only HEX, OPCODES, labels, builtins, and MACRO calls should be authorized.
     pub fn parse_body(&mut self) -> Result<Vec<Statement>, ParserError> {
         let mut statements: Vec<Statement> = Vec::new();
         self.match_kind(TokenKind::OpenBrace)?;
@@ -650,6 +650,21 @@ impl Parser {
                         span: AstSpan(vec![arg_span]),
                     });
                 }
+                TokenKind::BuiltinFunction(f) => {
+                    let mut curr_spans = vec![self.current_token.span.clone()];
+                    self.match_kind(TokenKind::BuiltinFunction(String::default()))?;
+                    let args = self.parse_args(true, false, false, true)?;
+                    args.iter().for_each(|a| curr_spans.extend_from_slice(&a.span.0));
+                    tracing::info!(target: "parser", "PARSING LABEL BODY: [BUILTIN FN: {}({:?})]", f, args);
+                    statements.push(Statement {
+                        ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
+                            kind: BuiltinFunctionKind::from(f),
+                            args,
+                            span: AstSpan(curr_spans.clone()),
+                        }),
+                        span: AstSpan(curr_spans),
+                    });
+                }
                 kind => {
                     tracing::error!(target: "parser", "TOKEN MISMATCH - LABEL BODY: {}", kind);
                     return Err(ParserError {


Test patch:
diff --git a/huff_parser/tests/labels.rs b/huff_parser/tests/labels.rs
index c73a86b..d87ca9d 100644
--- a/huff_parser/tests/labels.rs
+++ b/huff_parser/tests/labels.rs
@@ -175,3 +175,223 @@ fn multiline_labels() {
         assert_eq!(s.span, md_expected.statements[i].span);
     }
 }
+
+#[test]
+pub fn builtins_under_labels() {
+    let source = r#"
+    #define function myFunc() pure returns (uint256)
+
+    #define error TestError()
+
+    #define jumptable__packed TEST_TABLE {
+        my_label
+    }
+
+    #define macro SMALL_MACRO() = takes (3) returns (0) {
+        0x20 0x00 mstore
+    }
+
+    #define macro HELLO_WORLD() = takes(3) returns(0) {
+        my_label:
+            __tablestart(TEST_TABLE)
+            __tablesize(TEST_TABLE)
+            __codesize(SMALL_MACRO)
+            __FUNC_SIG(myFunc)
+            __ERROR(TestError)
+    }
+    "#;
+    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 macro_definition = parser.parse().unwrap().macros[1].clone();
+    let md_expected = MacroDefinition {
+        name: "HELLO_WORLD".to_string(),
+        parameters: vec![],
+        statements: vec![Statement {
+            ty: StatementType::Label(Label {
+                name: String::from("my_label"),
+                inner: vec![
+                    Statement {
+                        ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
+                            kind: BuiltinFunctionKind::Tablestart,
+                            args: vec![Argument {
+                                arg_type: None,
+                                name: Some(String::from("TEST_TABLE")),
+                                indexed: false,
+                                span: AstSpan(vec![Span { start: 342, end: 352, file: None }]),
+                            }],
+                            span: AstSpan(vec![
+                                Span { start: 329, end: 341, file: None },
+                                Span { start: 342, end: 352, file: None },
+                            ]),
+                        }),
+                        span: AstSpan(vec![
+                            Span { start: 329, end: 341, file: None },
+                            Span { start: 342, end: 352, file: None },
+                        ]),
+                    },
+                    Statement {
+                        ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
+                            kind: BuiltinFunctionKind::Tablesize,
+                            args: vec![Argument {
+                                arg_type: None,
+                                name: Some(String::from("TEST_TABLE")),
+                                indexed: false,
+                                span: AstSpan(vec![Span { start: 378, end: 388, file: None }]),
+                            }],
+                            span: AstSpan(vec![
+                                Span { start: 366, end: 377, file: None },
+                                Span { start: 378, end: 388, file: None },
+                            ]),
+                        }),
+                        span: AstSpan(vec![
+                            Span { start: 366, end: 377, file: None },
+                            Span { start: 378, end: 388, file: None },
+                        ]),
+                    },
+                    Statement {
+                        ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
+                            kind: BuiltinFunctionKind::Codesize,
+                            args: vec![Argument {
+                                arg_type: None,
+                                name: Some(String::from("SMALL_MACRO")),
+                                indexed: false,
+                                span: AstSpan(vec![Span { start: 413, end: 424, file: None }]),
+                            }],
+                            span: AstSpan(vec![
+                                Span { start: 402, end: 412, file: None },
+                                Span { start: 413, end: 424, file: None },
+                            ]),
+                        }),
+                        span: AstSpan(vec![
+                            Span { start: 402, end: 412, file: None },
+                            Span { start: 413, end: 424, file: None },
+                        ]),
+                    },
+                    Statement {
+                        ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
+                            kind: BuiltinFunctionKind::FunctionSignature,
+                            args: vec![Argument {
+                                arg_type: None,
+                                name: Some(String::from("myFunc")),
+                                indexed: false,
+                                span: AstSpan(vec![Span { start: 449, end: 455, file: None }]),
+                            }],
+                            span: AstSpan(vec![
+                                Span { start: 438, end: 448, file: None },
+                                Span { start: 449, end: 455, file: None },
+                            ]),
+                        }),
+                        span: AstSpan(vec![
+                            Span { start: 438, end: 448, file: None },
+                            Span { start: 449, end: 455, file: None },
+                        ]),
+                    },
+                    Statement {
+                        ty: StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
+                            kind: BuiltinFunctionKind::Error,
+                            args: vec![Argument {
+                                arg_type: None,
+                                name: Some(String::from("TestError")),
+                                indexed: false,
+                                span: AstSpan(vec![Span { start: 477, end: 486, file: None }]),
+                            }],
+                            span: AstSpan(vec![
+                                Span { start: 469, end: 476, file: None },
+                                Span { start: 477, end: 486, file: None },
+                            ]),
+                        }),
+                        span: AstSpan(vec![
+                            Span { start: 469, end: 476, file: None },
+                            Span { start: 477, end: 486, file: None },
+                        ]),
+                    },
+                ],
+                span: AstSpan(vec![
+                    Span { start: 307, end: 315, file: None },
+                    Span { start: 329, end: 341, file: None },
+                    Span { start: 342, end: 352, file: None },
+                    Span { start: 366, end: 377, file: None },
+                    Span { start: 378, end: 388, file: None },
+                    Span { start: 402, end: 412, file: None },
+                    Span { start: 413, end: 424, file: None },
+                    Span { start: 438, end: 448, file: None },
+                    Span { start: 449, end: 455, file: None },
+                    Span { start: 469, end: 476, file: None },
+                    Span { start: 477, end: 486, file: None },
+                ]),
+            }),
+            span: AstSpan(vec![
+                Span { start: 307, end: 315, file: None },
+                Span { start: 329, end: 341, file: None },
+                Span { start: 342, end: 352, file: None },
+                Span { start: 366, end: 377, file: None },
+                Span { start: 378, end: 388, file: None },
+                Span { start: 402, end: 412, file: None },
+                Span { start: 413, end: 424, file: None },
+                Span { start: 438, end: 448, file: None },
+                Span { start: 449, end: 455, file: None },
+                Span { start: 469, end: 476, file: None },
+                Span { start: 477, end: 486, file: None },
+            ]),
+        }],
+        takes: 3,
+        returns: 0,
+        span: AstSpan(vec![
+            Span { start: 247, end: 254, file: None },
+            Span { start: 255, end: 260, file: None },
+            Span { start: 261, end: 272, file: None },
+            Span { start: 272, end: 273, file: None },
+            Span { start: 273, end: 274, file: None },
+            Span { start: 275, end: 276, file: None },
+            Span { start: 277, end: 282, file: None },
+            Span { start: 282, end: 283, file: None },
+            Span { start: 283, end: 284, file: None },
+            Span { start: 284, end: 285, file: None },
+            Span { start: 286, end: 293, file: None },
+            Span { start: 293, end: 294, file: None },
+            Span { start: 294, end: 295, file: None },
+            Span { start: 295, end: 296, file: None },
+            Span { start: 297, end: 298, file: None },
+            Span { start: 307, end: 315, file: None },
+            Span { start: 315, end: 316, file: None },
+            Span { start: 329, end: 341, file: None },
+            Span { start: 341, end: 342, file: None },
+            Span { start: 342, end: 352, file: None },
+            Span { start: 352, end: 353, file: None },
+            Span { start: 366, end: 377, file: None },
+            Span { start: 377, end: 378, file: None },
+            Span { start: 378, end: 388, file: None },
+            Span { start: 388, end: 389, file: None },
+            Span { start: 402, end: 412, file: None },
+            Span { start: 412, end: 413, file: None },
+            Span { start: 413, end: 424, file: None },
+            Span { start: 424, end: 425, file: None },
+            Span { start: 438, end: 448, file: None },
+            Span { start: 448, end: 449, file: None },
+            Span { start: 449, end: 455, file: None },
+            Span { start: 455, end: 456, file: None },
+            Span { start: 469, end: 476, file: None },
+            Span { start: 476, end: 477, file: None },
+            Span { start: 477, end: 486, file: None },
+            Span { start: 486, end: 487, file: None },
+            Span { start: 492, end: 493, file: None },
+        ]),
+        outlined: false,
+    };
+    assert_eq!(macro_definition.name, md_expected.name);
+    assert_eq!(macro_definition.parameters, md_expected.parameters);
+    assert_eq!(macro_definition.takes, md_expected.takes);
+    assert_eq!(macro_definition.returns, md_expected.returns);
+    assert_eq!(parser.current_token.kind, TokenKind::Eof);
+    assert_eq!(macro_definition.span, md_expected.span);
+
+    // Test that each statement is the correct type
+    for (i, s) in macro_definition.statements.iter().enumerate() {
+        assert_eq!(s.ty, md_expected.statements[i].ty);
+        assert_eq!(s.span, md_expected.statements[i].span);
+    }
+}