Pass constant as parameter for buildin (#32)

cakevm · web-flow · commit 0fe9d8df27da · 2025-02-01T20:20:36.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,11 @@
 - Introduce new lexer token type `Bytes`, as not all hex should be parsed to bytes32.
   - For constants in code tables the bytes are copied as defined with e.g. leading zeros.
   - For all other cases leading zeros are removed and the smallest push operation is used.
+- New built-in function `__LEFTPAD` to a hex input or a function result in a code table.
+  - The function can only be used in a code table.
+  - Example: `__LEFTPAD(0x123)` -> `0000000000000000000000000000000000000000000000000000000000000123`
+- Allow to pass a constant as a parameter to `__RIGHTPAD` or `__LEFTPAD`.
+  - Example: `__RIGHTPAD([CONST])`
 
 ## [1.0.10] - 2025-01-31
 - Add windows binary to the release.
diff --git a/crates/codegen/src/irgen/builtin_function.rs b/crates/codegen/src/irgen/builtin_function.rs
@@ -1,3 +1,4 @@
+use crate::irgen::constants::constant_gen;
 use crate::Codegen;
 use alloy_primitives::{hex, keccak256};
 use huff_neo_utils::bytecode::{BytecodeRes, Bytes, CircularCodeSizeIndices, Jump, Jumps};
@@ -8,6 +9,7 @@ use huff_neo_utils::opcodes::Opcode;
 use huff_neo_utils::prelude::{
     BuiltinFunctionArg, BuiltinFunctionCall, BuiltinFunctionKind, Contract, MacroDefinition, MacroInvocation, TableDefinition,
 };
+use std::fmt::Display;
 
 // TODO: First step to refactor and split the function into smaller functions
 #[allow(clippy::too_many_arguments)]
@@ -98,10 +100,17 @@ pub fn builtin_function_gen<'a>(
             bytes.push((starting_offset, Bytes(push_bytes)));
         }
         BuiltinFunctionKind::RightPad => {
-            let push_bytes = right_pad(evm_version, contract, bf)?;
+            let push_bytes = builtin_pad(evm_version, contract, bf, PadDirection::Right)?;
             *offset += push_bytes.len() / 2;
             bytes.push((starting_offset, Bytes(push_bytes)));
         }
+        BuiltinFunctionKind::LeftPad => {
+            return Err(CodegenError {
+                kind: CodegenErrorKind::InvalidArguments(String::from("LeftPad is not supported in a function or macro")),
+                span: bf.span.clone(),
+                token: None,
+            });
+        }
         BuiltinFunctionKind::DynConstructorArg => {
             if bf.args.len() != 2 {
                 tracing::error!(
@@ -475,19 +484,38 @@ pub fn function_signature(contract: &Contract, bf: &BuiltinFunctionCall) -> Resu
     Ok(push_bytes)
 }
 
-pub fn right_pad(evm_version: &EVMVersion, contract: &Contract, bf: &BuiltinFunctionCall) -> Result<String, CodegenError> {
+#[derive(Debug, Clone, PartialEq)]
+pub enum PadDirection {
+    Left,
+    Right,
+}
+impl Display for PadDirection {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PadDirection::Left => write!(f, "__LEFTPAD"),
+            PadDirection::Right => write!(f, "__RIGHTPAD"),
+        }
+    }
+}
+
+pub fn builtin_pad(
+    evm_version: &EVMVersion,
+    contract: &Contract,
+    bf: &BuiltinFunctionCall,
+    direction: PadDirection,
+) -> Result<String, CodegenError> {
     if bf.args.len() != 1 {
-        tracing::error!(target = "codegen", "Incorrect number of arguments passed to __RIGHTPAD, should be 1: {}", bf.args.len());
+        tracing::error!(target = "codegen", "Incorrect number of arguments passed to {direction}, should be 1: {}", bf.args.len());
         return Err(CodegenError {
             kind: CodegenErrorKind::InvalidArguments(format!(
-                "Incorrect number of arguments passed to __RIGHTPAD, should be 1: {}",
+                "Incorrect number of arguments passed to {direction}, should be 1: {}",
                 bf.args.len()
             )),
             span: bf.span.clone(),
             token: None,
         });
     }
-    let first_arg = match bf.args[0] {
+    let first_arg = match &bf.args[0] {
         BuiltinFunctionArg::Argument(ref arg) => arg.name.clone().unwrap_or_default(),
         BuiltinFunctionArg::BuiltinFunctionCall(ref inner_call) => {
             match inner_call.kind {
@@ -500,27 +528,33 @@ pub fn right_pad(evm_version: &EVMVersion, contract: &Contract, bf: &BuiltinFunc
                     push_bytes[2..].to_string() // remove opcode
                 }
                 _ => {
-                    tracing::error!(target: "codegen", "Invalid argument type passed to __RIGHTPAD");
+                    tracing::error!(target: "codegen", "Invalid function call argument type passed to {direction}");
                     return Err(CodegenError {
-                        kind: CodegenErrorKind::InvalidArguments(String::from("Invalid argument type passed to __RIGHTPAD")),
+                        kind: CodegenErrorKind::InvalidArguments(format!("Invalid argument type passed to {direction}")),
                         span: bf.span.clone(),
                         token: None,
                     });
                 }
             }
         }
+        BuiltinFunctionArg::Constant(name, span) => {
+            let push_bytes = constant_gen(evm_version, name, contract, span)?;
+            push_bytes[2..].to_string() // remove opcode
+        }
         _ => {
-            tracing::error!(target: "codegen", "Invalid argument type passed to __RIGHTPAD");
+            tracing::error!(target: "codegen", "Invalid argument type passed to {direction}");
             return Err(CodegenError {
-                kind: CodegenErrorKind::InvalidArguments(String::from("Invalid argument type passed to __RIGHTPAD")),
+                kind: CodegenErrorKind::InvalidArguments(format!("Invalid argument type passed to {direction}")),
                 span: bf.span.clone(),
                 token: None,
             });
         }
     };
     let hex = format_even_bytes(first_arg);
-    let push_bytes = format!("{}{hex}{}", Opcode::Push32, "0".repeat(64 - hex.len()));
-    Ok(push_bytes)
+    if direction == PadDirection::Left {
+        return Ok(format!("{}{}{hex}", Opcode::Push32, "0".repeat(64 - hex.len())));
+    }
+    Ok(format!("{}{hex}{}", Opcode::Push32, "0".repeat(64 - hex.len())))
 }
 
 pub fn builtin_bytes(evm_version: &EVMVersion, bf: &BuiltinFunctionCall) -> Result<String, CodegenError> {
diff --git a/crates/codegen/src/lib.rs b/crates/codegen/src/lib.rs
@@ -3,7 +3,7 @@
 #![warn(unused_extern_crates)]
 #![forbid(unsafe_code)]
 
-use crate::irgen::builtin_function::{builtin_bytes, function_signature, right_pad};
+use crate::irgen::builtin_function::{builtin_bytes, builtin_pad, function_signature, PadDirection};
 use alloy_primitives::hex;
 use huff_neo_utils::ast::huff::*;
 use huff_neo_utils::ast::span::AstSpan;
@@ -163,7 +163,11 @@ impl Codegen {
                         byte_code = format!("{byte_code}{}", &bytes[2..]);
                     }
                     BuiltinFunctionKind::RightPad => {
-                        let bytes = right_pad(evm_version, contract, &bf)?;
+                        let bytes = builtin_pad(evm_version, contract, &bf, PadDirection::Right)?;
+                        byte_code = format!("{byte_code}{}", &bytes[2..]);
+                    }
+                    BuiltinFunctionKind::LeftPad => {
+                        let bytes = builtin_pad(evm_version, contract, &bf, PadDirection::Left)?;
                         byte_code = format!("{byte_code}{}", &bytes[2..]);
                     }
                     BuiltinFunctionKind::Bytes => {
diff --git a/crates/core/tests/code_table.rs b/crates/core/tests/code_table.rs
@@ -122,6 +122,44 @@ fn test_code_table_constant() {
     assert_eq!(mbytes, "60 14 61 0005 0101010101010101010101010101010101010101".replace(" ", ""));
 }
 
+#[test]
+fn test_code_table_leftpad() {
+    let source: &str = r#"
+        #define table CODE_TABLE() {
+            __LEFTPAD(0x0123)
+        }
+
+        #define macro MAIN() = takes(0) returns (0) {
+            __tablesize(CODE_TABLE)
+            __tablestart(CODE_TABLE)
+        }
+    "#;
+
+    // Parse tokens
+    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);
+
+    // Parse the AST
+    let mut contract = parser.parse().unwrap();
+
+    // Derive storage pointers
+    contract.derive_storage_pointers();
+
+    // Instantiate Codegen
+    let cg = Codegen::new();
+
+    // The codegen instance should have no artifact
+    assert!(cg.artifact.is_none());
+
+    // Have the Codegen create the constructor bytecode
+    let mbytes = Codegen::generate_main_bytecode(&EVMVersion::default(), &contract, None).unwrap();
+
+    // 60 = PUSH1, 20 = 32 bytes table size, 61 = PUSH2, 0005 = 5 bytes table start
+    assert_eq!(mbytes, "60 20 61 0005 0000000000000000000000000000000000000000000000000000000000000123".replace(" ", ""));
+}
+
 #[test]
 fn test_code_table_constant_with_leading_zeros() {
     let source: &str = r#"
@@ -207,6 +245,8 @@ fn test_code_table_all_features() {
             __RIGHTPAD(__FUNC_SIG("transfer(address,uint256)"))
             0xDEADBEEF
             [ADDR]
+            __LEFTPAD(0x123)
+            __LEFTPAD([ADDR])
         }
 
         #define macro MAIN() = takes(0) returns (0) {
@@ -236,10 +276,13 @@ fn test_code_table_all_features() {
     // Have the Codegen create the constructor bytecode
     let mbytes = Codegen::generate_main_bytecode(&EVMVersion::default(), &contract, None).unwrap();
 
-    // 60 = PUSH1, 36 = 56 bytes table size (32 + 4 + 20), 61 = PUSH2, 0005 = 5 bytes table start, a9059cbb = transfer(address,uint256)
+    // 60 = PUSH1, 78 = 120 bytes table size (32 + 4 + 20), 61 = PUSH2, 0005 = 5 bytes table start, a9059cbb = transfer(address,uint256)
     assert_eq!(
         mbytes,
-        "60 38 61 0005 a9059cbb00000000000000000000000000000000000000000000000000000000 deadbeef 0101010101010101010101010101010101010101"
+        "60 78 61 0005 a9059cbb00000000000000000000000000000000000000000000000000000000 deadbeef \
+        0101010101010101010101010101010101010101 \
+        0000000000000000000000000000000000000000000000000000000000000123\
+        0000000000000000000000000101010101010101010101010101010101010101"
             .replace(" ", "")
     );
 }
diff --git a/crates/parser/src/lib.rs b/crates/parser/src/lib.rs
@@ -807,49 +807,68 @@ impl Parser {
 
         tracing::debug!(target: "parser", "PARSING ARGs: {:?}", self.current_token.kind);
         while !self.check(TokenKind::CloseParen) {
-            // Check for strings
-            if let TokenKind::Str(s) = &self.current_token.kind {
-                args.push(BuiltinFunctionArg::Argument(Argument {
-                    name: Some(s.to_owned()), // Place the string in the "name" field
-                    arg_type: None,
-                    indexed: false,
-                    span: AstSpan(vec![self.current_token.span.clone()]),
-                    arg_location: None,
-                }));
-                self.consume();
-            }
-
-            // Check for literals
-            if let TokenKind::Literal(l) = &self.current_token.kind {
-                args.push(BuiltinFunctionArg::Argument(Argument {
-                    // Place literal in the "name" field
-                    name: Some(bytes32_to_hex_string(l, false)),
-                    arg_location: None,
-                    arg_type: None,
-                    indexed: false,
-                    span: AstSpan(vec![self.current_token.span.clone()]),
-                }));
-                self.consume();
-            }
-
-            // Check for identifiers
-            if let TokenKind::Ident(ident) = &self.current_token.kind {
-                args.push(BuiltinFunctionArg::Argument(Argument {
-                    name: Some(ident.to_owned()),
-                    arg_type: None,
-                    indexed: false,
-                    span: AstSpan(vec![self.current_token.span.clone()]),
-                    arg_location: None,
-                }));
-                self.consume();
-            }
-
             let mut is_builtin = false;
             let mut builtin = "".to_string();
-            if let TokenKind::BuiltinFunction(builtin_ref) = &self.current_token.kind {
-                is_builtin = true;
-                builtin = builtin_ref.clone();
-                self.consume();
+
+            match &self.current_token.kind {
+                TokenKind::Str(s) => {
+                    args.push(BuiltinFunctionArg::Argument(Argument {
+                        name: Some(s.to_owned()), // Place the string in the "name" field
+                        arg_type: None,
+                        indexed: false,
+                        span: AstSpan(vec![self.current_token.span.clone()]),
+                        arg_location: None,
+                    }));
+                    self.consume();
+                }
+                TokenKind::Bytes(bytes) => {
+                    args.push(BuiltinFunctionArg::Argument(Argument {
+                        name: Some(bytes.to_owned()),
+                        arg_type: None,
+                        indexed: false,
+                        span: AstSpan(vec![self.current_token.span.clone()]),
+                        arg_location: None,
+                    }));
+                    self.consume();
+                }
+                TokenKind::Literal(l) => {
+                    args.push(BuiltinFunctionArg::Argument(Argument {
+                        // Place literal in the "name" field
+                        name: Some(bytes32_to_hex_string(l, false)),
+                        arg_location: None,
+                        arg_type: None,
+                        indexed: false,
+                        span: AstSpan(vec![self.current_token.span.clone()]),
+                    }));
+                    self.consume();
+                }
+                TokenKind::Ident(ident) => {
+                    args.push(BuiltinFunctionArg::Argument(Argument {
+                        name: Some(ident.to_owned()),
+                        arg_type: None,
+                        indexed: false,
+                        span: AstSpan(vec![self.current_token.span.clone()]),
+                        arg_location: None,
+                    }));
+                    self.consume();
+                }
+                TokenKind::BuiltinFunction(builtin_ref) => {
+                    is_builtin = true;
+                    builtin = builtin_ref.clone();
+                    self.consume();
+                }
+                TokenKind::OpenBracket => {
+                    let (constant_name, span) = self.parse_constant_push()?;
+                    args.push(BuiltinFunctionArg::Constant(constant_name, AstSpan(vec![span])));
+                }
+                _ => {
+                    return Err(ParserError {
+                        kind: ParserErrorKind::UnexpectedType(self.current_token.kind.clone()),
+                        hint: Some("Expected string, literal, hex, identifier or build-in".to_string()),
+                        spans: AstSpan(vec![self.current_token.span.clone()]),
+                        cursor: self.cursor,
+                    });
+                }
             }
 
             if is_builtin {
@@ -1147,7 +1166,7 @@ impl Parser {
                         });
                     }
                     let mut curr_spans = vec![self.current_token.span.clone()];
-                    self.match_kind(TokenKind::BuiltinFunction(String::default()))?;
+                    self.consume();
                     let args = self.parse_builtin_args()?;
                     args.iter().for_each(|a| curr_spans.extend_from_slice(a.span().inner_ref()));
                     tracing::info!(target: "parser", "PARSING CODE TABLE BODY: [BUILTIN FN: {}({:?})]", f, args);
diff --git a/crates/utils/src/ast/huff.rs b/crates/utils/src/ast/huff.rs
@@ -533,6 +533,8 @@ pub enum BuiltinFunctionArg {
     BuiltinFunctionCall(BuiltinFunctionCall),
     /// Abi Argument
     Argument(Argument),
+    /// Constant Argument
+    Constant(String, AstSpan),
 }
 
 impl BuiltinFunctionArg {
@@ -542,6 +544,7 @@ impl BuiltinFunctionArg {
             BuiltinFunctionArg::Literal(_) => AstSpan::default(),
             BuiltinFunctionArg::BuiltinFunctionCall(b) => b.span.clone(),
             BuiltinFunctionArg::Argument(a) => a.span.clone(),
+            BuiltinFunctionArg::Constant(_, span) => span.clone(),
         }
     }
 }
@@ -577,6 +580,8 @@ pub enum BuiltinFunctionKind {
     Error,
     /// Rightpad function
     RightPad,
+    /// Leftpad function (only available in code tables)
+    LeftPad,
     /// Dynamic constructor arg function
     DynConstructorArg,
     /// Inject Raw Bytes
@@ -595,6 +600,7 @@ impl From<String> for BuiltinFunctionKind {
             "__EVENT_HASH" => BuiltinFunctionKind::EventHash,
             "__ERROR" => BuiltinFunctionKind::Error,
             "__RIGHTPAD" => BuiltinFunctionKind::RightPad,
+            "__LEFTPAD" => BuiltinFunctionKind::LeftPad,
             "__CODECOPY_DYN_ARG" => BuiltinFunctionKind::DynConstructorArg,
             "__VERBATIM" => BuiltinFunctionKind::Verbatim,
             "__BYTES" => BuiltinFunctionKind::Bytes,
@@ -617,6 +623,7 @@ impl TryFrom<&String> for BuiltinFunctionKind {
             "__EVENT_HASH" => Ok(BuiltinFunctionKind::EventHash),
             "__ERROR" => Ok(BuiltinFunctionKind::Error),
             "__RIGHTPAD" => Ok(BuiltinFunctionKind::RightPad),
+            "__LEFTPAD" => Ok(BuiltinFunctionKind::LeftPad),
             "__CODECOPY_DYN_ARG" => Ok(BuiltinFunctionKind::DynConstructorArg),
             "__VERBATIM" => Ok(BuiltinFunctionKind::Verbatim),
             "__BYTES" => Ok(BuiltinFunctionKind::Bytes),
@@ -635,6 +642,7 @@ impl Display for BuiltinFunctionKind {
             BuiltinFunctionKind::EventHash => write!(f, "__EVENT_HASH"),
             BuiltinFunctionKind::Error => write!(f, "__ERROR"),
             BuiltinFunctionKind::RightPad => write!(f, "__RIGHTPAD"),
+            BuiltinFunctionKind::LeftPad => write!(f, "__LEFTPAD"),
             BuiltinFunctionKind::DynConstructorArg => write!(f, "__CODECOPY_DYN_ARG"),
             BuiltinFunctionKind::Verbatim => write!(f, "__VERBATIM"),
             BuiltinFunctionKind::Bytes => write!(f, "__BYTES"),
