huff_neo_utils/ast/
huff.rs

1use crate::ast::abi::{Argument, EventDefinition, FunctionDefinition};
2use crate::ast::span::AstSpan;
3use crate::builtin_eval::{PadDirection, eval_builtin_bytes, eval_builtin_pad_simple, eval_event_hash, eval_function_signature};
4use crate::{
5    bytecode::*,
6    bytes_util::*,
7    error::{CodegenError, CodegenErrorKind},
8    evm_version::EVMVersion,
9    opcodes::{OPCODES_MAP, Opcode},
10    prelude::{MacroArg::Ident, Span, TokenKind},
11    push_value::PushValue,
12};
13use alloy_primitives::U256;
14use indexmap::IndexMap;
15use std::collections::HashSet;
16use std::{
17    collections::BTreeMap,
18    fmt::{Display, Formatter},
19    path::PathBuf,
20    sync::{Arc, Mutex},
21};
22
23/// A contained literal
24pub type Literal = [u8; 32];
25
26/// Reserved identifier accepted as the argument to `__codesize` to query the byte length of
27/// the runtime section (MAIN body + appended runtime tables). Disallowed as a user macro name.
28pub const RUNTIME_CODESIZE_ARG: &str = "RUNTIME";
29
30/// A File Path
31///
32/// Used for parsing the huff imports.
33pub type FilePath = PathBuf;
34
35/// A Huff Contract Representation
36///
37/// This is the representation of a contract as it is parsed from huff source code.
38/// Thus, it is also the root of the AST.
39///
40/// For examples of Huff contracts, see the [huff-examples repository](https://github.com/huff-language/huff-examples).
41#[derive(Debug, Default, Clone)]
42pub struct Contract {
43    // A IndexMap is used to preserve the insertion order of macros.
44    // This helps to keep the resulting bytecode deterministic.
45    /// Macro definitions, indexed by name
46    pub macros: IndexMap<String, MacroDefinition>,
47    /// Invocations of macros
48    pub invocations: Vec<MacroInvocation>,
49    /// File Imports
50    pub imports: Vec<FilePath>,
51    /// Constants
52    pub constants: Arc<Mutex<Vec<ConstantDefinition>>>,
53    /// Custom Errors
54    pub errors: Vec<ErrorDefinition>,
55    /// Functions
56    pub functions: Vec<FunctionDefinition>,
57    /// Events
58    pub events: Vec<EventDefinition>,
59    /// Tables
60    pub tables: Vec<TableDefinition>,
61    /// Labels
62    pub labels: HashSet<String>,
63    /// Flattened source code (for debugging and source maps)
64    pub flattened_source: Option<String>,
65    /// Individual source files for multi-file debugging support
66    /// Vec of (file_path, source_content)
67    pub source_files: Vec<(String, String)>,
68    /// Mapping from flattened position to (file_index, original_position)
69    pub source_map: Vec<(usize, usize, usize)>, // (flattened_start, file_id, original_start)
70    /// Byte length of the runtime section. Set before constructor codegen so
71    /// `__codesize(RUNTIME)` resolves; references from any other context error.
72    pub runtime_size: Option<usize>,
73}
74
75/// Returns true if `bf` is a `__codesize(RUNTIME)` call (single identifier arg `RUNTIME`).
76pub fn is_runtime_codesize(bf: &BuiltinFunctionCall) -> bool {
77    matches!(bf.kind, BuiltinFunctionKind::Codesize)
78        && matches!(bf.args.as_slice(), [BuiltinFunctionArg::Identifier(name, _)] if name == RUNTIME_CODESIZE_ARG)
79}
80
81impl Contract {
82    /// Map a position from the flattened source to the original file and position
83    /// Returns (file_id, original_start, original_end)
84    pub fn map_flattened_position_to_source(&self, flattened_start: usize, flattened_end: usize) -> (u32, usize, usize) {
85        // Find which file this position belongs to
86        let mut file_id = 0u32;
87        let mut original_start = flattened_start;
88        let mut original_end = flattened_end;
89
90        // Iterate through source_map to find the right file
91        for (idx, (file_start_in_flattened, fid, _)) in self.source_map.iter().enumerate() {
92            // Check if we have a next entry to determine the file's end position
93            let file_end = if idx + 1 < self.source_map.len() {
94                self.source_map[idx + 1].0
95            } else {
96                // Last file extends to the end
97                usize::MAX
98            };
99
100            // Check if this position is within this file's range
101            if *file_start_in_flattened <= flattened_start && flattened_start < file_end {
102                file_id = *fid as u32;
103                // Map to position within this file (subtract the file's start position in flattened source)
104                original_start = flattened_start - file_start_in_flattened;
105                original_end = flattened_end - file_start_in_flattened;
106                break;
107            }
108        }
109
110        (file_id, original_start, original_end)
111    }
112
113    /// Convert a span from original file coordinates to flattened source coordinates
114    /// This is needed because AST spans are in original file positions, but source maps need flattened positions
115    pub fn map_original_span_to_flattened(&self, span: &crate::prelude::Span) -> Option<(usize, usize)> {
116        if let Some(file) = &span.file {
117            // Find this file in our source_files list
118            let file_id = self.source_files.iter().enumerate().find(|(_, (path, _))| path == &file.path).map(|(id, _)| id)?;
119
120            // Find the start position of this file in the flattened source
121            let file_start_in_flattened = self.source_map.iter().find(|(_, fid, _)| *fid == file_id).map(|(start, _, _)| *start)?;
122
123            // Convert original file positions to flattened positions
124            let flattened_start = file_start_in_flattened + span.start;
125            let flattened_end = file_start_in_flattened + span.end;
126
127            Some((flattened_start, flattened_end))
128        } else {
129            None
130        }
131    }
132
133    /// Returns the first macro that matches the provided name
134    pub fn find_macro_by_name(&self, name: &str) -> Option<&MacroDefinition> {
135        self.macros.get(name).or_else(|| {
136            tracing::warn!("Failed to find macro \"{}\" in contract", name);
137            None
138        })
139    }
140
141    /// Returns the first table that matches the provided name
142    pub fn find_table_by_name(&self, name: &str) -> Option<TableDefinition> {
143        if let Some(t) = self.tables.iter().find(|t| t.name == name) {
144            Some(t.clone())
145        } else {
146            tracing::warn!("Failed to find table \"{}\" in contract", name);
147            None
148        }
149    }
150
151    /// Derives the FreeStoragePointers into their bytes32 representation
152    ///
153    /// Slots are assigned in order of first use: the constructor, then MAIN, then each test with its
154    /// decorators, so pointers only tests use never move the contract's own slots.
155    pub fn derive_storage_pointers(&mut self) {
156        let mut storage_pointers: Vec<(String, [u8; 32])> = Vec::new();
157        let mut last_assigned_free_pointer = 0;
158
159        // Derive Constructor Storage Pointers
160        match self.find_macro_by_name("CONSTRUCTOR") {
161            Some(m) => self.recurse_ast_constants(m, &mut storage_pointers, &mut last_assigned_free_pointer, false),
162            None => {
163                // The constructor is not required, so we can just warn
164                tracing::warn!(target: "ast", "'CONSTRUCTOR' MACRO NOT FOUND WHILE DERIVING STORAGE POINTERS!")
165            }
166        }
167
168        // Derive Main Storage Pointers
169        match self.find_macro_by_name("MAIN") {
170            Some(m) => self.recurse_ast_constants(m, &mut storage_pointers, &mut last_assigned_free_pointer, false),
171            None => {
172                tracing::error!(target: "ast", "'MAIN' MACRO NOT FOUND WHILE DERIVING STORAGE POINTERS!")
173            }
174        }
175
176        // Derive pointers that only tests use, in definition order. They come after the contract's own
177        // pointers, so tests never change the slots of the deployed contract.
178        for test in self.macros.values().filter(|m| m.test) {
179            self.recurse_ast_constants(test, &mut storage_pointers, &mut last_assigned_free_pointer, false);
180            let decorator_constants =
181                test.decorator.iter().flat_map(|d| &d.flags).flat_map(|f| &f.args).filter_map(|arg| match &arg.value {
182                    DecoratorValue::Constant(name) => Some(name),
183                    _ => None,
184                });
185            for name in decorator_constants {
186                self.assign_free_storage_pointers(name, &test.name, &mut storage_pointers, &mut last_assigned_free_pointer);
187            }
188        }
189
190        tracing::debug!(target: "ast", "Generate Storage pointers: {:?}", storage_pointers);
191        tracing::debug!(target: "ast", "ALL AST CONSTANTS: {:?}", storage_pointers);
192
193        // Set all the constants to their new values
194        for c in self.constants.lock().unwrap().iter_mut() {
195            match storage_pointers.iter().filter(|pointer| pointer.0.eq(&c.name)).collect::<Vec<&(String, [u8; 32])>>().first() {
196                Some(p) => {
197                    *c = ConstantDefinition { name: c.name.to_string(), value: ConstVal::StoragePointer(p.1), span: c.span.clone() };
198                }
199                None => {
200                    tracing::warn!(target: "ast", "SET STORAGE POINTER BUT FAILED TO SET DERIVED CONSTANT VALUE FOR \"{}\"", c.name)
201                }
202            }
203        }
204    }
205
206    /// Recurse down an AST Macro Definition to set Storage Pointers
207    ///
208    /// ## Overview
209    ///
210    /// For each statement in the macro definition:
211    ///     - If it's a free storage pointer constant, set the constant value if not already set and
212    ///       updated out `last_p` tracker value
213    ///     - If it's a literal constant, we can set the constant value directly to the literal if
214    ///       not already set
215    ///     - If it's a macro invocation, look for the macro definition and recurse into that macro
216    ///       definition using `recurse_ast_constants`
217    pub fn recurse_ast_constants(
218        &self,
219        macro_def: &MacroDefinition,
220        storage_pointers: &mut Vec<(String, [u8; 32])>,
221        last_p: &mut i32,
222        checking_constructor: bool,
223    ) {
224        let mut visited = std::collections::HashSet::new();
225        self.recurse_ast_constants_inner(macro_def, storage_pointers, last_p, checking_constructor, &mut visited)
226    }
227
228    fn recurse_ast_constants_inner(
229        &self,
230        macro_def: &MacroDefinition,
231        storage_pointers: &mut Vec<(String, [u8; 32])>,
232        last_p: &mut i32,
233        checking_constructor: bool,
234        visited: &mut std::collections::HashSet<String>,
235    ) {
236        // Check for circular recursion
237        if visited.contains(&macro_def.name) {
238            tracing::warn!(target: "ast", "Circular macro invocation detected: '{}' is already being processed. Skipping to prevent infinite recursion.", macro_def.name);
239            return;
240        }
241
242        // Mark this macro as being processed
243        visited.insert(macro_def.name.clone());
244        let mut statements = macro_def.statements.clone();
245
246        let mut i = 0;
247        loop {
248            if i >= statements.len() {
249                break;
250            }
251            match &statements[i].clone().ty {
252                StatementType::Constant(const_name) => {
253                    self.assign_free_storage_pointers(const_name, &macro_def.name, storage_pointers, last_p);
254                }
255                StatementType::MacroInvocation(mi) => {
256                    tracing::debug!(target: "ast", "Found macro invocation: \"{}\" in macro def: \"{}\"!", mi.macro_name, macro_def.name);
257
258                    // Check for constant references in macro arguments
259                    let mut constant_args: Vec<String> = Vec::new();
260                    for arg in &mi.args {
261                        // check if it is a constant
262                        if let Ident(name) = arg {
263                            self.constants.lock().unwrap().iter().for_each(|constant| {
264                                if name == &constant.name {
265                                    tracing::debug!(target: "ast", "CONSTANT FOUND AS MACRO PARAMETER {}", name);
266                                    constant_args.push(name.to_string());
267                                }
268                            })
269                        }
270                    }
271                    // Assign constants that reference the Free Storage Pointer
272                    for constant_arg in constant_args {
273                        self.assign_free_storage_pointers(&constant_arg, &macro_def.name, storage_pointers, last_p);
274                    }
275
276                    match self.macros.get(&mi.macro_name) {
277                        Some(md) => {
278                            if md.name.eq("CONSTRUCTOR") {
279                                if !checking_constructor {
280                                    self.recurse_ast_constants_inner(md, storage_pointers, last_p, true, visited);
281                                }
282                            } else {
283                                self.recurse_ast_constants_inner(md, storage_pointers, last_p, checking_constructor, visited);
284                            }
285                        }
286                        None => {
287                            tracing::warn!(target: "ast", "MACRO \"{}\" INVOKED BUT NOT FOUND IN AST!", mi.macro_name)
288                        }
289                    }
290                }
291                StatementType::BuiltinFunctionCall(bfc) => {
292                    tracing::debug!(target: "ast", "Deriving Storage Pointers: Found builtin function {:?}", bfc.kind);
293                    for builtin_fn_arg in &bfc.args {
294                        let BuiltinFunctionArg::Identifier(name, _) = builtin_fn_arg else { continue };
295                        match self.macros.get(name) {
296                            Some(md) => {
297                                if md.name.eq("CONSTRUCTOR") {
298                                    if !checking_constructor {
299                                        self.recurse_ast_constants_inner(md, storage_pointers, last_p, true, visited);
300                                    }
301                                } else {
302                                    self.recurse_ast_constants_inner(md, storage_pointers, last_p, checking_constructor, visited);
303                                }
304                            }
305                            None => {
306                                tracing::warn!(target: "ast", "BUILTIN HAS ARG NAME \"{}\" BUT NOT FOUND IN AST!", name)
307                            }
308                        }
309                    }
310                }
311                StatementType::Label(l) => {
312                    for state in l.inner.iter().rev() {
313                        statements.insert(i + 1, state.clone());
314                    }
315                }
316                _ => {}
317            }
318            i += 1;
319        }
320
321        // Remove this macro from visited set when done processing it
322        visited.remove(&macro_def.name);
323
324        // Breadth-first
325        // if !macros_to_recurse.is_empty() {
326        //     let next_md = macros_to_recurse.remove(0);
327        //     self.recurse_ast_constants(next_md, storage_pointers, last_p, macros_to_recurse);
328        // }
329    }
330
331    fn assign_free_storage_pointers(
332        &self,
333        const_name: &String,
334        macro_name: &String,
335        storage_pointers: &mut Vec<(String, [u8; 32])>,
336        last_p: &mut i32,
337    ) {
338        tracing::debug!(target: "ast", "Found constant \"{}\" in macro def \"{}\" statements!", const_name, macro_name);
339        if storage_pointers.iter().filter(|pointer| pointer.0.eq(const_name)).collect::<Vec<&(String, [u8; 32])>>().is_empty() {
340            tracing::debug!(target: "ast", "No storage pointer already set for \"{}\"!", const_name);
341            // Get the associated constant
342            match self.constants.lock().unwrap().iter().filter(|c| c.name.eq(const_name)).collect::<Vec<&ConstantDefinition>>().first() {
343                Some(c) => {
344                    match c.value {
345                        ConstVal::FreeStoragePointer(_) => {
346                            let old_p = *last_p;
347                            *last_p += 1;
348                            let new_value = str_to_bytes32(&format!("{old_p}"));
349                            storage_pointers.push((const_name.to_string(), new_value));
350                        }
351                        ConstVal::Bytes(_)
352                        | ConstVal::String(_)
353                        | ConstVal::BuiltinFunctionCall(_)
354                        | ConstVal::Expression(_)
355                        | ConstVal::Noop => {
356                            // Skip constants that are not free storage pointers
357                        }
358                        // This should never be reached, as we only assign free storage pointers
359                        _ => panic!("Invalid Constant Value"),
360                    };
361                }
362                None => {
363                    tracing::warn!(target: "ast", "CONSTANT \"{}\" NOT FOUND IN AST CONSTANTS", const_name)
364                }
365            }
366        }
367    }
368
369    /// Add override constants to the AST
370    ///
371    /// ## Overview
372    ///
373    /// For each override constant, add it to the AST if it doesn't already exist. Override
374    /// constants can be passed in via the CLI.
375    pub fn add_override_constants(&self, override_constants: &Option<BTreeMap<&str, Bytes>>) {
376        if let Some(override_constants) = override_constants {
377            for (name, value) in override_constants {
378                let mut constants = self.constants.lock().unwrap();
379                if let Some(c) = constants.iter_mut().find(|c| c.name.as_str().eq(*name)) {
380                    c.value = ConstVal::Bytes(value.clone());
381                } else {
382                    constants.push(ConstantDefinition {
383                        name: name.to_string(),
384                        value: ConstVal::Bytes(value.clone()),
385                        span: AstSpan::default(),
386                    });
387                }
388            }
389        }
390    }
391
392    /// Validate that all opcodes used in the contract are compatible with the target EVM version
393    pub fn validate_opcodes(&self, evm_version: &EVMVersion) -> Result<(), CodegenError> {
394        for macro_def in self.macros.values() {
395            Self::validate_statements_opcodes(&macro_def.statements, evm_version)?;
396        }
397
398        for table in &self.tables {
399            Self::validate_statements_opcodes(&table.statements, evm_version)?;
400        }
401
402        Ok(())
403    }
404
405    /// Helper function to recursively validate opcodes in a list of statements
406    fn validate_statements_opcodes(statements: &[Statement], evm_version: &EVMVersion) -> Result<(), CodegenError> {
407        for statement in statements {
408            match &statement.ty {
409                StatementType::Opcode(opcode) | StatementType::ImmediateOpcode { opcode, .. } => {
410                    if let Some(required_version) = opcode.requires_evm_version()
411                        && evm_version.version() < &required_version
412                    {
413                        return Err(CodegenError {
414                            kind: CodegenErrorKind::InvalidOpcodeForEVMVersion(
415                                opcode.as_ref().to_string(),
416                                required_version.to_string(),
417                                evm_version.to_string(),
418                            ),
419                            span: statement.span.clone_box(),
420                            token: Some(TokenKind::Opcode(*opcode)),
421                        });
422                    }
423                }
424                StatementType::Label(label) => {
425                    // Recursively validate statements inside labels
426                    Self::validate_statements_opcodes(&label.inner, evm_version)?;
427                }
428                StatementType::ForLoop { body, .. } => {
429                    Self::validate_statements_opcodes(body, evm_version)?;
430                }
431                StatementType::IfStatement { then_branch, else_if_branches, else_branch, .. } => {
432                    Self::validate_statements_opcodes(then_branch, evm_version)?;
433                    for (_, branch) in else_if_branches {
434                        Self::validate_statements_opcodes(branch, evm_version)?;
435                    }
436                    if let Some(branch) = else_branch {
437                        Self::validate_statements_opcodes(branch, evm_version)?;
438                    }
439                }
440                StatementType::MacroInvocation(invocation) => {
441                    // Validate opcodes passed as arguments to the macro
442                    Self::validate_macro_invocation_args(&invocation.args, evm_version, &statement.span)?;
443                }
444                StatementType::ArgMacroInvocation(_, _, args) => {
445                    // Validate opcodes passed as arguments to the arg macro invocation
446                    Self::validate_macro_invocation_args(args, evm_version, &statement.span)?;
447                }
448                _ => {
449                    // Other statement types don't contain opcodes we need to validate
450                }
451            }
452        }
453        Ok(())
454    }
455
456    /// Helper function to validate opcodes passed as macro invocation arguments
457    fn validate_macro_invocation_args(args: &[MacroArg], evm_version: &EVMVersion, invocation_span: &AstSpan) -> Result<(), CodegenError> {
458        for arg in args {
459            match arg {
460                MacroArg::Ident(name) => {
461                    // Check if this identifier is an opcode name
462                    if let Some(&opcode) = OPCODES_MAP.get(name.as_str()) {
463                        // Validate the opcode against the EVM version
464                        if let Some(required_version) = opcode.requires_evm_version()
465                            && evm_version.version() < &required_version
466                        {
467                            return Err(CodegenError {
468                                kind: CodegenErrorKind::InvalidOpcodeForEVMVersion(
469                                    opcode.as_ref().to_string(),
470                                    required_version.to_string(),
471                                    evm_version.to_string(),
472                                ),
473                                span: invocation_span.clone_box(),
474                                token: Some(TokenKind::Opcode(opcode)),
475                            });
476                        }
477                    }
478                }
479                MacroArg::MacroCall(invocation) => {
480                    // Recursively validate nested macro invocations
481                    Self::validate_macro_invocation_args(&invocation.args, evm_version, &invocation.span)?;
482                }
483                MacroArg::ArgCallMacroInvocation(_, nested_args) => {
484                    // Recursively validate arg call macro invocations
485                    // Note: We use the parent invocation span since arg calls don't have their own span
486                    Self::validate_macro_invocation_args(nested_args, evm_version, invocation_span)?;
487                }
488                _ => {
489                    // Literals and ArgCalls don't need validation
490                }
491            }
492        }
493        Ok(())
494    }
495
496    /// Recursively evaluates arithmetic expressions in constant definitions at compile time.
497    /// This supports binary operations (+, -, *, /, %), unary negation, literals, constant references,
498    /// and parenthesized expressions with proper operator precedence and overflow/underflow checking.
499    pub fn evaluate_constant_expression(&self, expr: &Expression) -> Result<Literal, CodegenError> {
500        // Start evaluation with an empty set of currently-evaluating constants
501        let mut evaluating = std::collections::HashSet::new();
502        self.evaluate_constant_expression_internal(expr, &mut evaluating)
503    }
504
505    /// Internal helper that tracks which constants are currently being evaluated to detect cycles
506    fn evaluate_constant_expression_internal(
507        &self,
508        expr: &Expression,
509        evaluating: &mut std::collections::HashSet<String>,
510    ) -> Result<Literal, CodegenError> {
511        match expr {
512            Expression::Literal { value, .. } => {
513                // Direct literal value
514                Ok(*value)
515            }
516
517            Expression::Constant { name, span } => {
518                // Check for circular dependency
519                if evaluating.contains(name) {
520                    return Err(CodegenError {
521                        kind: CodegenErrorKind::CircularConstantDependency(name.clone()),
522                        span: span.clone_box(),
523                        token: None,
524                    });
525                }
526
527                // Mark this constant as being evaluated
528                evaluating.insert(name.clone());
529
530                // Look up the constant by name and clone the value
531                // We must clone before recursing to avoid holding the lock during recursion
532                let const_value = {
533                    let constants = self.constants.lock().unwrap();
534                    let const_def = constants.iter().find(|c| c.name == *name).ok_or_else(|| CodegenError {
535                        kind: CodegenErrorKind::UndefinedConstant(name.clone()),
536                        span: span.clone_box(),
537                        token: None,
538                    })?;
539                    const_def.value.clone()
540                };
541
542                // Evaluate the constant's value (lock is released at this point)
543                let result = match &const_value {
544                    ConstVal::Bytes(bytes) => {
545                        // Parse hex string to [u8; 32]
546                        Ok(str_to_bytes32(&bytes.as_str()))
547                    }
548                    ConstVal::Expression(e) => {
549                        // Recursively evaluate nested expressions
550                        self.evaluate_constant_expression_internal(e, evaluating)
551                    }
552                    ConstVal::StoragePointer(lit) => {
553                        // Storage pointers are already resolved literals
554                        Ok(*lit)
555                    }
556                    ConstVal::BuiltinFunctionCall(bf) => {
557                        // Evaluate builtin functions at compile time
558                        let push_value = match bf.kind {
559                            BuiltinFunctionKind::FunctionSignature => eval_function_signature(self, bf)?,
560                            BuiltinFunctionKind::EventHash => eval_event_hash(self, bf)?,
561                            BuiltinFunctionKind::Bytes => eval_builtin_bytes(bf)?,
562                            BuiltinFunctionKind::RightPad => eval_builtin_pad_simple(bf, PadDirection::Right)?,
563                            BuiltinFunctionKind::LeftPad => eval_builtin_pad_simple(bf, PadDirection::Left)?,
564                            // __codesize(RUNTIME) — other __codesize(MACRO) calls depend on
565                            // layout details unavailable in constant context.
566                            _ if is_runtime_codesize(bf) => {
567                                let size = self.runtime_size.ok_or_else(|| CodegenError {
568                                    kind: CodegenErrorKind::RuntimeSizeNotComputed,
569                                    span: span.clone_box(),
570                                    token: None,
571                                })?;
572                                PushValue::from(U256::from(size).to_be_bytes::<32>())
573                            }
574                            _ => {
575                                return Err(CodegenError {
576                                    kind: CodegenErrorKind::UnsupportedBuiltinFunction(format!("{}", bf.kind)),
577                                    span: span.clone_box(),
578                                    token: None,
579                                });
580                            }
581                        };
582                        // Extract the [u8; 32] value from PushValue
583                        Ok(push_value.value.0)
584                    }
585                    _ => Err(CodegenError { kind: CodegenErrorKind::InvalidConstantExpression, span: span.clone_box(), token: None }),
586                };
587
588                // Remove from evaluating set after evaluation completes
589                evaluating.remove(name);
590
591                result
592            }
593
594            Expression::Binary { left, op, right, span } => {
595                // Evaluate both operands
596                let left_val = self.evaluate_constant_expression_internal(left, evaluating)?;
597                let right_val = self.evaluate_constant_expression_internal(right, evaluating)?;
598
599                // Convert to U256 for arithmetic
600                let l = U256::from_be_bytes(left_val);
601                let r = U256::from_be_bytes(right_val);
602
603                // Perform the operation with overflow checking
604                let result = match op {
605                    BinaryOp::Add => l.checked_add(r).ok_or_else(|| CodegenError {
606                        kind: CodegenErrorKind::ArithmeticOverflow,
607                        span: span.clone_box(),
608                        token: None,
609                    })?,
610                    BinaryOp::Sub => l.checked_sub(r).ok_or_else(|| CodegenError {
611                        kind: CodegenErrorKind::ArithmeticUnderflow,
612                        span: span.clone_box(),
613                        token: None,
614                    })?,
615                    BinaryOp::Mul => l.checked_mul(r).ok_or_else(|| CodegenError {
616                        kind: CodegenErrorKind::ArithmeticOverflow,
617                        span: span.clone_box(),
618                        token: None,
619                    })?,
620                    BinaryOp::Div => {
621                        if r.is_zero() {
622                            return Err(CodegenError { kind: CodegenErrorKind::DivisionByZero, span: span.clone_box(), token: None });
623                        }
624                        l / r
625                    }
626                    BinaryOp::Mod => {
627                        if r.is_zero() {
628                            return Err(CodegenError { kind: CodegenErrorKind::DivisionByZero, span: span.clone_box(), token: None });
629                        }
630                        l % r
631                    }
632                    // Comparison operators return 1 (true) or 0 (false)
633                    BinaryOp::Eq => {
634                        if l == r {
635                            U256::from(1u64)
636                        } else {
637                            U256::ZERO
638                        }
639                    }
640                    BinaryOp::NotEq => {
641                        if l != r {
642                            U256::from(1u64)
643                        } else {
644                            U256::ZERO
645                        }
646                    }
647                    BinaryOp::Lt => {
648                        if l < r {
649                            U256::from(1u64)
650                        } else {
651                            U256::ZERO
652                        }
653                    }
654                    BinaryOp::Gt => {
655                        if l > r {
656                            U256::from(1u64)
657                        } else {
658                            U256::ZERO
659                        }
660                    }
661                    BinaryOp::LtEq => {
662                        if l <= r {
663                            U256::from(1u64)
664                        } else {
665                            U256::ZERO
666                        }
667                    }
668                    BinaryOp::GtEq => {
669                        if l >= r {
670                            U256::from(1u64)
671                        } else {
672                            U256::ZERO
673                        }
674                    }
675                };
676
677                Ok(result.to_be_bytes())
678            }
679
680            Expression::Unary { op, expr, .. } => {
681                // Evaluate the operand
682                let val = self.evaluate_constant_expression_internal(expr, evaluating)?;
683                let v = U256::from_be_bytes(val);
684
685                // Perform unary operation
686                let result = match op {
687                    UnaryOp::Neg => {
688                        // Two's complement negation
689                        U256::ZERO.wrapping_sub(v)
690                    }
691                    UnaryOp::Not => {
692                        // Logical NOT: 0 becomes 1, non-zero becomes 0
693                        if v.is_zero() { U256::from(1u64) } else { U256::ZERO }
694                    }
695                };
696
697                Ok(result.to_be_bytes())
698            }
699
700            Expression::Grouped { expr, .. } => {
701                // Grouped expressions just evaluate the inner expression
702                self.evaluate_constant_expression_internal(expr, evaluating)
703            }
704
705            Expression::ArgCall { name, span, .. } => {
706                // ArgCall cannot be evaluated without macro invocation context
707                // This should only be reached during global expansion (before macro expansion)
708                // When if/for expansion moves to statement_gen, this won't be called
709                Err(CodegenError { kind: CodegenErrorKind::ArgCallInConstantExpression(name.clone()), span: span.clone_box(), token: None })
710            }
711        }
712    }
713}
714
715/// A Table Definition
716#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
717pub struct TableDefinition {
718    /// The name of the table
719    pub name: String,
720    /// The table kind
721    pub kind: TableKind,
722    /// The table's statements
723    pub statements: Vec<Statement>,
724    /// Size of table if known. If built-in function is used, this will be `None` and the size will be calculated during codegen.
725    pub size: Option<Literal>,
726    /// The table span
727    pub span: AstSpan,
728}
729
730impl TableDefinition {
731    /// Public associated function that instantiates a TableDefinition from a string
732    pub fn new(name: String, kind: TableKind, statements: Vec<Statement>, size: Option<Literal>, span: AstSpan) -> Self {
733        TableDefinition { name, kind, statements, size, span }
734    }
735}
736
737/// A Table Kind
738#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
739pub enum TableKind {
740    /// A regular jump table
741    JumpTable,
742    /// A packed jump table
743    JumpTablePacked,
744    /// A code table
745    CodeTable,
746}
747
748impl From<TokenKind> for TableKind {
749    /// Public associated function that converts a TokenKind to a TableKind
750    fn from(token_kind: TokenKind) -> Self {
751        match token_kind {
752            TokenKind::JumpTable => TableKind::JumpTable,
753            TokenKind::JumpTablePacked => TableKind::JumpTablePacked,
754            TokenKind::CodeTable => TableKind::CodeTable,
755            _ => panic!("Invalid Token Kind"), // TODO: Better error handling
756        }
757    }
758}
759
760/// A Macro Definition
761#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
762pub struct MacroDefinition {
763    /// The Macro Name
764    pub name: String,
765    /// The macro's decorator
766    pub decorator: Option<Decorator>,
767    /// A list of Macro parameters
768    pub parameters: Vec<Argument>,
769    /// A list of Statements contained in the Macro
770    pub statements: Vec<Statement>,
771    /// The take size
772    pub takes: usize,
773    /// The return size
774    pub returns: usize,
775    /// The Span of the Macro Definition
776    pub span: AstSpan,
777    /// Is the macro a function (outlined)?
778    pub outlined: bool,
779    /// Is the macro a test?
780    pub test: bool,
781}
782
783impl ToIRBytecode<CodegenError> for MacroDefinition {
784    fn to_irbytecode(&self, evm_version: &EVMVersion) -> Result<IRBytecode<'_>, CodegenError> {
785        let inner_irbytes: Vec<IRBytes> = MacroDefinition::to_irbytes(evm_version, &self.statements);
786        Ok(IRBytecode(inner_irbytes))
787    }
788}
789
790impl MacroDefinition {
791    /// Public associated function that instantiates a MacroDefinition.
792    #[allow(clippy::too_many_arguments)]
793    pub fn new(
794        name: String,
795        decorator: Option<Decorator>,
796        parameters: Vec<Argument>,
797        statements: Vec<Statement>,
798        takes: usize,
799        returns: usize,
800        spans: Vec<Span>,
801        outlined: bool,
802        test: bool,
803    ) -> Self {
804        MacroDefinition { name, decorator, parameters, statements, takes, returns, span: AstSpan(spans), outlined, test }
805    }
806
807    /// Translate statements into IRBytes
808    pub fn to_irbytes<'a>(evm_version: &EVMVersion, statements: &'a [Statement]) -> Vec<IRBytes<'a>> {
809        let mut inner_irbytes: Vec<IRBytes> = vec![];
810
811        let mut statement_iter = statements.iter();
812        while let Some(statement) = statement_iter.next() {
813            match &statement.ty {
814                StatementType::Literal(l) => {
815                    let push_bytes = literal_gen(evm_version, l);
816                    inner_irbytes.push(IRBytes { ty: IRByteType::Bytes(Bytes::Raw(push_bytes)), span: &statement.span });
817                }
818                StatementType::Opcode(o) => {
819                    let opcode_str = o.string();
820                    inner_irbytes.push(IRBytes { ty: IRByteType::Bytes(Bytes::Raw(opcode_str)), span: &statement.span });
821                    // If the opcode is a push that takes a literal value, we need to consume the
822                    // next statement, which must be a literal as checked in the parser
823                    if o.is_value_push() {
824                        match statement_iter.next() {
825                            Some(Statement { ty: StatementType::Literal(l), span: _ }) => {
826                                let hex_literal: String = bytes32_to_hex_string(l, false);
827                                let prefixed_hex_literal = o.prefix_push_literal(&hex_literal);
828                                inner_irbytes
829                                    .push(IRBytes { ty: IRByteType::Bytes(Bytes::Raw(prefixed_hex_literal)), span: &statement.span });
830                            }
831                            _ => {
832                                // We have a push without a literal - this should be caught by the
833                                // parser
834                                panic!("Invalid push statement");
835                            }
836                        }
837                    }
838                }
839                StatementType::Code(c) => {
840                    inner_irbytes.push(IRBytes { ty: IRByteType::Bytes(Bytes::Raw(c.to_owned())), span: &statement.span });
841                }
842                StatementType::MacroInvocation(mi) => {
843                    inner_irbytes.push(IRBytes {
844                        ty: IRByteType::Statement(Box::new(Statement {
845                            ty: StatementType::MacroInvocation(mi.clone()),
846                            span: statement.span.clone(),
847                        })),
848                        span: &statement.span,
849                    });
850                }
851                StatementType::Constant(name) => {
852                    // Constant needs to be evaluated at the top-level
853                    inner_irbytes.push(IRBytes { ty: IRByteType::Constant(name.to_owned()), span: &statement.span });
854                }
855                StatementType::ArgCall(parent_macro_name, arg_name) => {
856                    // Arg call needs to use a destination defined in the calling macro context
857                    inner_irbytes.push(IRBytes {
858                        ty: IRByteType::ArgCall(parent_macro_name.to_owned(), arg_name.to_owned()),
859                        span: &statement.span,
860                    });
861                }
862                StatementType::ArgMacroInvocation(parent_macro_name, arg_name, args) => {
863                    // Arg macro invocation - will be resolved to actual macro invocation during codegen
864                    inner_irbytes.push(IRBytes {
865                        ty: IRByteType::Statement(Box::new(Statement {
866                            ty: StatementType::ArgMacroInvocation(parent_macro_name.clone(), arg_name.clone(), args.clone()),
867                            span: statement.span.clone(),
868                        })),
869                        span: &statement.span,
870                    });
871                }
872                StatementType::LabelCall(jump_to) => {
873                    /* Jump To doesn't translate directly to bytecode */
874                    inner_irbytes.push(IRBytes {
875                        ty: IRByteType::Statement(Box::new(Statement {
876                            ty: StatementType::LabelCall(jump_to.to_string()),
877                            span: statement.span.clone(),
878                        })),
879                        span: &statement.span,
880                    });
881                }
882                StatementType::Label(l) => {
883                    /* Jump Dests don't translate directly to bytecode */
884                    inner_irbytes.push(IRBytes {
885                        ty: IRByteType::Statement(Box::new(Statement {
886                            ty: StatementType::Label(l.clone()),
887                            span: statement.span.clone(),
888                        })),
889                        span: &statement.span,
890                    });
891
892                    // Recurse label statements to IRBytes Bytes
893                    inner_irbytes.append(&mut MacroDefinition::to_irbytes(evm_version, &l.inner));
894                }
895                StatementType::BuiltinFunctionCall(builtin) => {
896                    inner_irbytes.push(IRBytes {
897                        ty: IRByteType::Statement(Box::new(Statement {
898                            ty: StatementType::BuiltinFunctionCall(builtin.clone()),
899                            span: statement.span.clone(),
900                        })),
901                        span: &statement.span,
902                    });
903                }
904                StatementType::ImmediateOpcode { .. } => {
905                    // Operands may reference constants, macro arguments or loop variables, so they are
906                    // resolved during statement_gen
907                    inner_irbytes.push(IRBytes { ty: IRByteType::Statement(Box::new(statement.clone())), span: &statement.span });
908                }
909                StatementType::ForLoop { .. } => {
910                    // ForLoop will be expanded during statement_gen with macro invocation context
911                    inner_irbytes.push(IRBytes { ty: IRByteType::Statement(Box::new(statement.clone())), span: &statement.span });
912                }
913                StatementType::IfStatement { .. } => {
914                    // IfStatement will be expanded during statement_gen with macro invocation context
915                    inner_irbytes.push(IRBytes { ty: IRByteType::Statement(Box::new(statement.clone())), span: &statement.span });
916                }
917            }
918        }
919
920        inner_irbytes
921    }
922}
923
924/// A Macro Invocation
925#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
926pub struct MacroInvocation {
927    /// The Macro Name
928    pub macro_name: String,
929    /// A list of Macro arguments
930    pub args: Vec<MacroArg>,
931    /// The Macro Invocation Span
932    pub span: AstSpan,
933}
934
935/// An argument passed when invoking a macro
936#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
937pub enum MacroArg {
938    /// Hexadecimal literal argument
939    ///
940    /// Example: `MACRO(0x04)`, `MACRO(0x420)`
941    HexLiteral(Literal),
942    /// Macro Iden String Argument
943    ///
944    /// Example: `MACRO(my_label)`, `MACRO(SOME_CONSTANT)`
945    Ident(String),
946    /// An Arg Call
947    ///
948    /// Passing through an argument from outer macro. Example: `INNER(<value>)`
949    ArgCall(ArgCall),
950    /// An Arg Call that is invoked as a macro
951    ///
952    /// First-class macro invocation. Example: `APPLY(<operation>())`
953    ArgCallMacroInvocation(String, Vec<MacroArg>),
954    /// A Nested Macro Call
955    ///
956    /// Example: `OUTER(INNER())`, `TRANSFER(CALCULATE_AMOUNT())`
957    MacroCall(MacroInvocation),
958    /// Opcode Argument
959    ///
960    /// Example: `APPLY_OP(add)`, `USE_TWO_OPS(tload, clz)`
961    Opcode(Opcode),
962    /// Builtin Function Call Argument
963    ///
964    /// Example: `MACRO(__FUNC_SIG(transfer))`, `MACRO(__EVENT_HASH(Transfer))`
965    BuiltinFunctionCall(BuiltinFunctionCall),
966    /// __NOOP builtin constant (generates no bytecode)
967    ///
968    /// Example: `MACRO(__NOOP)`
969    Noop,
970}
971
972/// An argument call
973#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
974pub struct ArgCall {
975    /// The Macro Name
976    pub macro_name: String,
977    /// The name of the argument
978    pub name: String,
979    /// The span of the argument call
980    pub span: AstSpan,
981}
982
983/// Free Storage Pointer Unit Struct
984#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
985pub struct FreeStoragePointer;
986
987/// An arithmetic expression node
988#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
989pub enum Expression {
990    /// A literal value
991    Literal {
992        /// The literal value
993        value: Literal,
994        /// The span of the literal
995        span: AstSpan,
996    },
997    /// Reference to a named constant
998    Constant {
999        /// The name of the constant
1000        name: String,
1001        /// The span of the constant reference
1002        span: AstSpan,
1003    },
1004    /// Binary operation (e.g., 1 + 2, 3 * 4)
1005    Binary {
1006        /// Left operand
1007        left: Box<Expression>,
1008        /// Binary operator
1009        op: BinaryOp,
1010        /// Right operand
1011        right: Box<Expression>,
1012        /// The span of the binary operation
1013        span: AstSpan,
1014    },
1015    /// Unary operation (e.g., -5)
1016    Unary {
1017        /// Unary operator
1018        op: UnaryOp,
1019        /// Operand expression
1020        expr: Box<Expression>,
1021        /// The span of the unary operation
1022        span: AstSpan,
1023    },
1024    /// Grouped/parenthesized expression
1025    Grouped {
1026        /// The inner expression
1027        expr: Box<Expression>,
1028        /// The span of the grouped expression (including parentheses)
1029        span: AstSpan,
1030    },
1031    /// Reference to a macro argument (e.g., `<arg>`)
1032    ArgCall {
1033        /// The parent macro name
1034        macro_name: String,
1035        /// The argument name
1036        name: String,
1037        /// The span of the arg call
1038        span: AstSpan,
1039    },
1040}
1041
1042impl Expression {
1043    /// Get the span of this expression
1044    pub fn span(&self) -> &AstSpan {
1045        match self {
1046            Expression::Literal { span, .. } => span,
1047            Expression::Constant { span, .. } => span,
1048            Expression::Binary { span, .. } => span,
1049            Expression::Unary { span, .. } => span,
1050            Expression::Grouped { span, .. } => span,
1051            Expression::ArgCall { span, .. } => span,
1052        }
1053    }
1054}
1055
1056/// Binary operators for arithmetic and comparison expressions
1057#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1058pub enum BinaryOp {
1059    /// Addition (+)
1060    Add,
1061    /// Subtraction (-)
1062    Sub,
1063    /// Multiplication (*)
1064    Mul,
1065    /// Division (/)
1066    Div,
1067    /// Modulo (%)
1068    Mod,
1069    /// Equal comparison (==)
1070    Eq,
1071    /// Not equal comparison (!=)
1072    NotEq,
1073    /// Less than (<)
1074    Lt,
1075    /// Greater than (>)
1076    Gt,
1077    /// Less than or equal (<=)
1078    LtEq,
1079    /// Greater than or equal (>=)
1080    GtEq,
1081}
1082
1083/// Unary operators for arithmetic and logical expressions
1084#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1085pub enum UnaryOp {
1086    /// Negation (-)
1087    Neg,
1088    /// Logical NOT (!)
1089    Not,
1090}
1091
1092/// A Constant Value
1093#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1094pub enum ConstVal {
1095    /// Bytes value for the constant
1096    Bytes(Bytes),
1097    /// String literal value for the constant
1098    String(String),
1099    /// A Free Storage Pointer
1100    FreeStoragePointer(FreeStoragePointer),
1101    /// A Storage Pointer assigned by the compiler
1102    StoragePointer(Literal),
1103    /// Built-in function call
1104    BuiltinFunctionCall(BuiltinFunctionCall),
1105    /// An arithmetic expression (evaluated at compile time)
1106    Expression(Expression),
1107    /// __NOOP builtin constant (generates no bytecode)
1108    Noop,
1109}
1110
1111/// A Constant Definition
1112#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1113pub struct ConstantDefinition {
1114    /// The Constant name
1115    pub name: String,
1116    /// The Constant value
1117    pub value: ConstVal,
1118    /// The Span of the Constant Definition
1119    pub span: AstSpan,
1120}
1121
1122/// An Error Definition
1123#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1124pub struct ErrorDefinition {
1125    /// The Error name
1126    pub name: String,
1127    /// The Error's selector
1128    pub selector: [u8; 4],
1129    /// The parameters of the error
1130    pub parameters: Vec<Argument>,
1131    /// The Span of the Constant Definition
1132    pub span: AstSpan,
1133}
1134
1135/// A Jump Destination
1136#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1137pub struct Label {
1138    /// The JumpDest Name
1139    pub name: String,
1140    /// Statements Inside The JumpDest
1141    pub inner: Vec<Statement>,
1142    /// The label span
1143    pub span: AstSpan,
1144}
1145
1146/// A Builtin Function Argument
1147#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1148pub enum BuiltinFunctionArg {
1149    /// A String Literal (quoted string)
1150    StringLiteral(String, AstSpan),
1151    /// An Identifier (macro name, table name, function name)
1152    Identifier(String, AstSpan),
1153    /// A Hex Literal (raw hex bytes)
1154    HexLiteral(String, AstSpan),
1155    /// Builtin Function Call
1156    BuiltinFunctionCall(BuiltinFunctionCall),
1157    /// Constant Argument
1158    Constant(String, AstSpan),
1159    /// A compile-time expression, e.g. `<arg>` or `([BASE] + 0x20)`
1160    Expression(Expression),
1161}
1162
1163impl BuiltinFunctionArg {
1164    /// Get the span of the Builtin Function Argument
1165    pub fn span(&self) -> AstSpan {
1166        match self {
1167            BuiltinFunctionArg::StringLiteral(_, span) => span.clone(),
1168            BuiltinFunctionArg::Identifier(_, span) => span.clone(),
1169            BuiltinFunctionArg::HexLiteral(_, span) => span.clone(),
1170            BuiltinFunctionArg::BuiltinFunctionCall(b) => b.span.clone(),
1171            BuiltinFunctionArg::Constant(_, span) => span.clone(),
1172            BuiltinFunctionArg::Expression(expr) => expr.span().clone(),
1173        }
1174    }
1175}
1176
1177/// A Builtin Function Call
1178#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1179pub struct BuiltinFunctionCall {
1180    /// The Builtin Kind
1181    pub kind: BuiltinFunctionKind,
1182    /// Arguments for the builtin function call.
1183    /// TODO: Maybe make a better type for this other than `Argument`? Would be nice if it pointed
1184    ///       directly to the macro/table.
1185    /// Update: Spitted with BuiltinFunctionArg, but still needs to be update for macro/table
1186    pub args: Vec<BuiltinFunctionArg>,
1187    /// The builtin function call span
1188    pub span: AstSpan,
1189}
1190
1191/// A Builtin Function Kind
1192#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1193pub enum BuiltinFunctionKind {
1194    /// Table size function
1195    Tablesize,
1196    /// Code size function
1197    Codesize,
1198    /// Table start function
1199    Tablestart,
1200    /// Embed table function
1201    EmbedTable,
1202    /// Function signature function
1203    FunctionSignature,
1204    /// Event hash function
1205    EventHash,
1206    /// Error selector function
1207    Error,
1208    /// Rightpad function
1209    RightPad,
1210    /// Leftpad function (only available in code tables)
1211    LeftPad,
1212    /// Offset in the deployment bytecode where the constructor arguments start
1213    ArgsStart,
1214    /// Inject Raw Bytes
1215    Verbatim,
1216    /// Bytes function to convert a string to bytes
1217    Bytes,
1218    /// Assert PC (program counter) position
1219    AssertPc,
1220}
1221
1222impl From<String> for BuiltinFunctionKind {
1223    fn from(value: String) -> Self {
1224        match value.as_str() {
1225            "__tablesize" => BuiltinFunctionKind::Tablesize,
1226            "__codesize" => BuiltinFunctionKind::Codesize,
1227            "__tablestart" => BuiltinFunctionKind::Tablestart,
1228            "__EMBED_TABLE" => BuiltinFunctionKind::EmbedTable,
1229            "__FUNC_SIG" => BuiltinFunctionKind::FunctionSignature,
1230            "__EVENT_HASH" => BuiltinFunctionKind::EventHash,
1231            "__ERROR" => BuiltinFunctionKind::Error,
1232            "__RIGHTPAD" => BuiltinFunctionKind::RightPad,
1233            "__LEFTPAD" => BuiltinFunctionKind::LeftPad,
1234            "__argsstart" => BuiltinFunctionKind::ArgsStart,
1235            "__VERBATIM" => BuiltinFunctionKind::Verbatim,
1236            "__BYTES" => BuiltinFunctionKind::Bytes,
1237            "__ASSERT_PC" => BuiltinFunctionKind::AssertPc,
1238            _ => panic!("Invalid Builtin Function Kind"), /* This should never be reached,
1239                                                           * builtins are validated with a
1240                                                           * `try_from` call in the lexer. */
1241        }
1242    }
1243}
1244
1245impl TryFrom<&String> for BuiltinFunctionKind {
1246    type Error = ();
1247
1248    fn try_from(value: &String) -> Result<Self, <BuiltinFunctionKind as TryFrom<&String>>::Error> {
1249        match value.as_str() {
1250            "__tablesize" => Ok(BuiltinFunctionKind::Tablesize),
1251            "__codesize" => Ok(BuiltinFunctionKind::Codesize),
1252            "__tablestart" => Ok(BuiltinFunctionKind::Tablestart),
1253            "__EMBED_TABLE" => Ok(BuiltinFunctionKind::EmbedTable),
1254            "__FUNC_SIG" => Ok(BuiltinFunctionKind::FunctionSignature),
1255            "__EVENT_HASH" => Ok(BuiltinFunctionKind::EventHash),
1256            "__ERROR" => Ok(BuiltinFunctionKind::Error),
1257            "__RIGHTPAD" => Ok(BuiltinFunctionKind::RightPad),
1258            "__LEFTPAD" => Ok(BuiltinFunctionKind::LeftPad),
1259            "__argsstart" => Ok(BuiltinFunctionKind::ArgsStart),
1260            "__VERBATIM" => Ok(BuiltinFunctionKind::Verbatim),
1261            "__BYTES" => Ok(BuiltinFunctionKind::Bytes),
1262            "__ASSERT_PC" => Ok(BuiltinFunctionKind::AssertPc),
1263            _ => Err(()),
1264        }
1265    }
1266}
1267
1268impl Display for BuiltinFunctionKind {
1269    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1270        match self {
1271            BuiltinFunctionKind::Tablesize => write!(f, "__tablesize"),
1272            BuiltinFunctionKind::Codesize => write!(f, "__codesize"),
1273            BuiltinFunctionKind::Tablestart => write!(f, "__tablestart"),
1274            BuiltinFunctionKind::EmbedTable => write!(f, "__EMBED_TABLE"),
1275            BuiltinFunctionKind::FunctionSignature => write!(f, "__FUNC_SIG"),
1276            BuiltinFunctionKind::EventHash => write!(f, "__EVENT_HASH"),
1277            BuiltinFunctionKind::Error => write!(f, "__ERROR"),
1278            BuiltinFunctionKind::RightPad => write!(f, "__RIGHTPAD"),
1279            BuiltinFunctionKind::LeftPad => write!(f, "__LEFTPAD"),
1280            BuiltinFunctionKind::ArgsStart => write!(f, "__argsstart"),
1281            BuiltinFunctionKind::Verbatim => write!(f, "__VERBATIM"),
1282            BuiltinFunctionKind::Bytes => write!(f, "__BYTES"),
1283            BuiltinFunctionKind::AssertPc => write!(f, "__ASSERT_PC"),
1284        }
1285    }
1286}
1287
1288impl BuiltinFunctionKind {
1289    /// Checks if a string matches any reserved builtin function name
1290    pub fn is_reserved_name(name: &str) -> bool {
1291        matches!(
1292            name,
1293            "__tablesize"
1294                | "__codesize"
1295                | "__tablestart"
1296                | "__EMBED_TABLE"
1297                | "__FUNC_SIG"
1298                | "__EVENT_HASH"
1299                | "__ERROR"
1300                | "__RIGHTPAD"
1301                | "__LEFTPAD"
1302                | "__argsstart"
1303                | "__VERBATIM"
1304                | "__BYTES"
1305                | "__ASSERT_PC"
1306                | "__NOOP"
1307        )
1308    }
1309}
1310
1311/// A Statement
1312#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1313pub struct Statement {
1314    /// The type of statement
1315    pub ty: StatementType,
1316    /// The span of the Statement
1317    pub span: AstSpan,
1318}
1319
1320/// The Statement Type
1321#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1322pub enum StatementType {
1323    /// A Literal Statement
1324    ///
1325    /// Example: `0x04`, `0x420`
1326    Literal(Literal),
1327    /// An Opcode Statement
1328    ///
1329    /// Example: `add`, `mstore`, `calldataload`
1330    Opcode(Opcode),
1331    /// An opcode followed by immediate operands that are resolved at compile time
1332    ///
1333    /// Operands are compile-time expressions: literals, `[CONSTANT]`, `<arg>`, loop variables, or
1334    /// parenthesized arithmetic. Example: `dupn 17`, `swapn [DEPTH]`, `exchange 1 <m>`, `push2 [SIZE]`
1335    ImmediateOpcode {
1336        /// The opcode
1337        opcode: Opcode,
1338        /// Operand expressions in source order
1339        operands: Vec<Expression>,
1340    },
1341    /// A Code Statement
1342    ///
1343    /// Raw bytecode string
1344    Code(String),
1345    /// A Macro Invocation Statement
1346    ///
1347    /// Example: `TRANSFER()`, `ONLY_OWNER()`
1348    MacroInvocation(MacroInvocation),
1349    /// A Constant Push
1350    ///
1351    /// Example: `[BALANCE_LOCATION]`, `[MAX_SUPPLY]`
1352    Constant(String),
1353    /// An Arg Call
1354    ///
1355    /// Referencing a macro parameter. Example: `<amount>`, `<recipient>`
1356    ///
1357    /// Macro name and argument name
1358    ArgCall(String, String),
1359    /// A Macro Invocation through Argument
1360    ///
1361    /// Calling a macro passed as an argument. Example: `<operation>()`
1362    ///
1363    /// Parent macro name, argument name, and arguments for the invoked macro
1364    ArgMacroInvocation(String, String, Vec<MacroArg>),
1365    /// A Label
1366    ///
1367    /// Jump destination definition. Example: `success:`, `error:`
1368    Label(Label),
1369    /// A Label Reference/Call
1370    ///
1371    /// Reference to jump to a label. Example: `success`, `error`
1372    LabelCall(String),
1373    /// A built-in function call
1374    ///
1375    /// Example: `__FUNC_SIG(transfer)`, `__EVENT_HASH(Transfer)`
1376    BuiltinFunctionCall(BuiltinFunctionCall),
1377    /// A compile-time for loop
1378    ///
1379    /// Example: `for(i in 0..5) { <i> }`
1380    ForLoop {
1381        /// Loop variable name
1382        variable: String,
1383        /// Start value (constant expression)
1384        start: Expression,
1385        /// End value (exclusive, constant expression)
1386        end: Expression,
1387        /// Optional step value (defaults to 1)
1388        step: Option<Expression>,
1389        /// Loop body statements
1390        body: Vec<Statement>,
1391    },
1392    /// A compile-time if statement
1393    ///
1394    /// Example: `if([CONST] > 5) { ... } else { ... }`
1395    IfStatement {
1396        /// Condition (constant expression that evaluates to non-zero = true)
1397        condition: Expression,
1398        /// Then branch statements
1399        then_branch: Vec<Statement>,
1400        /// Optional else if clauses (condition, body pairs)
1401        else_if_branches: Vec<(Expression, Vec<Statement>)>,
1402        /// Optional else branch statements
1403        else_branch: Option<Vec<Statement>>,
1404    },
1405}
1406
1407impl Display for StatementType {
1408    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1409        match self {
1410            StatementType::Literal(l) => write!(f, "LITERAL: {}", bytes32_to_hex_string(l, true)),
1411            StatementType::Opcode(o) => write!(f, "OPCODE: {o}"),
1412            StatementType::ImmediateOpcode { opcode, operands } => {
1413                write!(f, "IMMEDIATE OPCODE: {opcode:?} with {} operands", operands.len())
1414            }
1415            StatementType::Code(s) => write!(f, "CODE: {s}"),
1416            StatementType::MacroInvocation(m) => {
1417                write!(f, "MACRO INVOCATION: {}", m.macro_name)
1418            }
1419            StatementType::Constant(c) => write!(f, "CONSTANT: {c}"),
1420            StatementType::ArgCall(m, c) => write!(f, "ARG CALL in {m}: {c}"),
1421            StatementType::ArgMacroInvocation(parent_macro, arg_name, args) => {
1422                write!(f, "ARG MACRO INVOCATION: <{}>({:?}) in {}", arg_name, args, parent_macro)
1423            }
1424            StatementType::Label(l) => write!(f, "LABEL: {}", l.name),
1425            StatementType::LabelCall(l) => write!(f, "LABEL CALL: {l}"),
1426            StatementType::BuiltinFunctionCall(b) => {
1427                write!(f, "BUILTIN FUNCTION CALL: {:?}", b.kind)
1428            }
1429            StatementType::ForLoop { variable, start: _, end: _, step, body } => {
1430                write!(f, "FOR LOOP: {} in ... step {:?}, {} statements", variable, step, body.len())
1431            }
1432            StatementType::IfStatement { condition: _, then_branch, else_if_branches, else_branch } => {
1433                write!(
1434                    f,
1435                    "IF STATEMENT: then {} statements, {} else if branches, else branch: {}",
1436                    then_branch.len(),
1437                    else_if_branches.len(),
1438                    else_branch.is_some()
1439                )
1440            }
1441        }
1442    }
1443}
1444
1445/// A decorator tag
1446///
1447/// One or more `#[...]` blocks placed over a test definition. Developers use decorators to set the
1448/// environment, the starting state and the expected outcome of an individual test.
1449#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1450pub struct Decorator {
1451    /// Flags of every decorator block, merged in source order
1452    pub flags: Vec<DecoratorFlag>,
1453}
1454
1455/// A decorator flag such as `calldata("0x01")` or `store([SLOT], 0x01, at = 0x1000)`
1456///
1457/// The parser accepts any flag name; the test runner decides which names and arguments are valid.
1458#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1459pub struct DecoratorFlag {
1460    /// The flag name
1461    pub name: String,
1462    /// The arguments in source order; empty when the flag has no parentheses
1463    pub args: Vec<DecoratorArg>,
1464    /// The flag span
1465    pub span: AstSpan,
1466}
1467
1468/// A single decorator flag argument, optionally named (`max = 30000`)
1469#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1470pub struct DecoratorArg {
1471    /// The argument name for `key = value` arguments
1472    pub key: Option<String>,
1473    /// The argument value
1474    pub value: DecoratorValue,
1475    /// The argument span
1476    pub span: AstSpan,
1477}
1478
1479/// The value of a decorator flag argument
1480#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
1481pub enum DecoratorValue {
1482    /// A hex literal without the `0x` prefix, keeping its written width
1483    Hex(String),
1484    /// A decimal literal
1485    Int(usize),
1486    /// A string literal
1487    Str(String),
1488    /// `true` or `false`
1489    Bool(bool),
1490    /// A constant reference `[NAME]`
1491    Constant(String),
1492    /// A builtin function call such as `__FUNC_SIG(transfer)`
1493    Builtin(BuiltinFunctionCall),
1494    /// A bare identifier: a label, macro or error name, or a keyword such as `this`
1495    Ident(String),
1496}
huff_neo_utils/ast/huff.rs

huff_neo_utils/ast/
huff.rs
