🐛 Standalone nonpayable & FREE_STORAGE_POINTER (#14)

* 🐛 Fix function and constant parsing * 🥢 Fix clippy
huff-language · Oct 18, 2024 · d184772 · d184772
1 parent 83b3dec
commit d184772
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 13 deletions.
diff --git a/crates/ast/src/ast.rs b/crates/ast/src/ast.rs
@@ -17,7 +17,7 @@ pub enum Definition<'src> {
     Macro(Macro<'src>),
     Constant {
         name: Spanned<&'src str>,
-        value: U256,
+        expr: Spanned<ConstExpr>,
     },
     Jumptable(Jumptable<'src>),
     Table {
@@ -37,6 +37,12 @@ pub struct Macro<'src> {
     pub body: Box<[MacroStatement<'src>]>,
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub enum ConstExpr {
+    Value(U256),
+    FreeStoragePointer,
+}
+
 #[derive(Debug, PartialEq, Eq)]
 pub enum MacroStatement<'src> {
     LabelDefinition(Spanned<&'src str>),
@@ -69,7 +75,7 @@ pub enum Invoke<'src> {
 
 #[derive(Debug, PartialEq, Eq)]
 pub struct Jumptable<'src> {
-    pub name: (Span, &'src str),
+    pub name: Spanned<&'src str>,
     pub size: u8,
     pub labels: Box<[&'src str]>,
 }

diff --git a/crates/ast/src/parser.rs b/crates/ast/src/parser.rs
@@ -259,11 +259,19 @@ fn invoke<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::Invoke<'s
 }
 
 fn constant<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::Definition<'src>> {
+    let const_expr = choice((
+        word().map(|(value, span)| (ast::ConstExpr::Value(value), span)),
+        just(Ident("FREE_STORAGE_POINTER"))
+            .ignore_then(just(Punct('(')))
+            .ignore_then(just(Punct(')')))
+            .map_with(|_, ex| (ast::ConstExpr::FreeStoragePointer, ex.span())),
+    ));
+
     just(Ident("constant"))
         .ignore_then(ident())
         .then_ignore(just(Punct('=')))
-        .then(word())
-        .map(|(name, (value, _))| ast::Definition::Constant { name, value })
+        .then(const_expr)
+        .map(|(name, expr)| ast::Definition::Constant { name, expr })
 }
 
 fn table<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::Definition<'src>> {
@@ -287,13 +295,18 @@ fn sol_function<'tokens, 'src: 'tokens>() -> impl Parser<'tokens, 'src, ast::Def
         .ignore_then(ident())
         .then(sol_type_list())
         .then(
-            choice((just(Ident("public")), just(Ident("external"))))
-                .then_ignore(choice((just(Ident("view")), just(Ident("pure")))).or_not())
-                .then_ignore(choice((just(Ident("payable")), just(Ident("nonpayable")))).or_not())
-                .or_not()
-                .ignore_then(just(Ident("returns")))
-                .ignore_then(sol_type_list())
-                .or_not(),
+            choice((
+                just(Ident("public")),
+                just(Ident("external")),
+                just(Ident("payable")),
+                just(Ident("nonpayable")),
+            ))
+            .or_not()
+            .then_ignore(choice((just(Ident("view")), just(Ident("pure")))).or_not())
+            .or_not()
+            .ignore_then(just(Ident("returns")))
+            .ignore_then(sol_type_list())
+            .or_not(),
         )
         .map(|((name, args), rets)| {
             let rets = rets.unwrap_or(Box::new([]));
@@ -614,15 +627,36 @@ mod tests {
     }
 
     #[test]
-    fn parse_constant() {
+    fn parse_constant_value() {
         let span: Span = SimpleSpan::new(0, 0);
 
         assert_ok!(
             constant(),
             vec![Ident("constant"), Ident("TEST"), Punct('='), Hex("0x1")],
             ast::Definition::Constant {
                 name: ("TEST", span),
-                value: uint!(1_U256)
+                expr: (ast::ConstExpr::Value(uint!(1_U256)), span)
+            }
+        );
+    }
+
+    #[test]
+    fn parse_constant_storage_pointer() {
+        let span: Span = SimpleSpan::new(0, 0);
+
+        assert_ok!(
+            constant(),
+            vec![
+                Ident("constant"),
+                Ident("VAR_LOCATION"),
+                Punct('='),
+                Ident("FREE_STORAGE_POINTER"),
+                Punct('('),
+                Punct(')')
+            ],
+            ast::Definition::Constant {
+                name: ("VAR_LOCATION", span),
+                expr: (ast::ConstExpr::FreeStoragePointer, span)
             }
         );
     }