huff_neo_codegen/irgen/
statements.rs

1use alloy_primitives::U256;
2use huff_neo_utils::bytecode::{BytecodeSegments, JumpPlaceholderData, PushOpcode};
3use huff_neo_utils::bytes_util::str_to_bytes32;
4use huff_neo_utils::prelude::*;
5use huff_neo_utils::scope::ScopeManager;
6use std::collections::BTreeMap;
7
8use crate::Codegen;
9use crate::irgen::arg_calls::bubble_arg_call;
10use crate::irgen::builtin_function;
11use crate::irgen::constants::{constant_gen, lookup_constant};
12
13/// Type alias for statement generation result: (bytes, spans)
14type StatementGenResult = (BytecodeSegments, Vec<Option<(usize, usize)>>);
15
16/// Evaluates an expression with access to macro invocation context for resolving `<arg>` calls
17pub(crate) fn evaluate_expression_with_context<'a>(
18    expr: &Expression,
19    contract: &Contract,
20    scope_mgr: &ScopeManager<'a>,
21) -> Result<[u8; 32], CodegenError> {
22    match expr {
23        Expression::ArgCall { name, span: arg_span, .. } => {
24            // Search through the macro invocation stack for the argument's value
25            let mis = scope_mgr.invocation_vec();
26            for (_, parent_mi) in mis.iter().rev() {
27                if let Some(parent_macro) = contract.find_macro_by_name(&parent_mi.macro_name)
28                    && let Some(param_idx) = parent_macro.parameters.iter().position(|p| p.name.as_deref() == Some(name))
29                    && let Some(arg_value) = parent_mi.args.get(param_idx)
30                {
31                    // Found the argument - now evaluate it
32                    match arg_value {
33                        MacroArg::HexLiteral(lit) => {
34                            // Direct hex literal value
35                            return Ok(*lit);
36                        }
37                        MacroArg::Ident(ident_name) => {
38                            // Identifier - could be a constant reference, try to look it up and evaluate
39                            if let Ok(const_def) = lookup_constant(ident_name, contract, arg_span) {
40                                match &const_def.value {
41                                    ConstVal::Bytes(bytes) => {
42                                        use huff_neo_utils::bytes_util::str_to_bytes32;
43                                        return Ok(str_to_bytes32(&bytes.as_str()));
44                                    }
45                                    ConstVal::String(_s) => {
46                                        return Err(CodegenError {
47                                            kind: CodegenErrorKind::StringConstantNotBytes(format!(
48                                                "String constant cannot be used in expressions. Use __BYTES([{}]) to convert it to bytes.",
49                                                ident_name
50                                            )),
51                                            span: arg_span.clone_box(),
52                                            token: None,
53                                        });
54                                    }
55                                    ConstVal::Expression(const_expr) => {
56                                        return contract.evaluate_constant_expression(const_expr);
57                                    }
58                                    _ => {
59                                        return Err(CodegenError {
60                                            kind: CodegenErrorKind::InvalidConstantExpression,
61                                            span: arg_span.clone_box(),
62                                            token: None,
63                                        });
64                                    }
65                                }
66                            } else {
67                                // Not a constant - this is an error in expression context
68                                return Err(CodegenError {
69                                    kind: CodegenErrorKind::InvalidMacroArgumentType(format!(
70                                        "Argument '{}' resolved to identifier '{}' which cannot be used in expression",
71                                        name, ident_name
72                                    )),
73                                    span: arg_span.clone_box(),
74                                    token: None,
75                                });
76                            }
77                        }
78                        MacroArg::BuiltinFunctionCall(bf) => {
79                            // Builtin function call - evaluate it
80                            let push_value = Codegen::gen_builtin_bytecode(contract, bf, bf.span.clone())?;
81                            return Ok(push_value.value.0);
82                        }
83                        MacroArg::ArgCall(nested_arg) => {
84                            // Nested argument reference - recursively resolve
85                            return evaluate_expression_with_context(
86                                &Expression::ArgCall {
87                                    macro_name: String::new(),
88                                    name: nested_arg.name.clone(),
89                                    span: nested_arg.span.clone(),
90                                },
91                                contract,
92                                scope_mgr,
93                            );
94                        }
95                        _ => {
96                            return Err(CodegenError {
97                                kind: CodegenErrorKind::InvalidMacroArgumentType(format!(
98                                    "Argument '{}' cannot be used in expression (unsupported type)",
99                                    name
100                                )),
101                                span: arg_span.clone_box(),
102                                token: None,
103                            });
104                        }
105                    }
106                }
107            }
108
109            // Argument not found in invocation stack
110            Err(CodegenError { kind: CodegenErrorKind::MissingArgumentDefinition(name.clone()), span: arg_span.clone_box(), token: None })
111        }
112
113        Expression::Binary { left, op, right, span } => {
114            // Evaluate both operands using this function to handle nested ArgCalls
115            let left_val = evaluate_expression_with_context(left, contract, scope_mgr)?;
116            let right_val = evaluate_expression_with_context(right, contract, scope_mgr)?;
117
118            // Convert to U256 for arithmetic
119            let l = U256::from_be_bytes(left_val);
120            let r = U256::from_be_bytes(right_val);
121
122            // Perform the operation
123            let result = match op {
124                BinaryOp::Add => l.wrapping_add(r),
125                BinaryOp::Sub => l.wrapping_sub(r),
126                BinaryOp::Mul => l.wrapping_mul(r),
127                BinaryOp::Div => {
128                    if r.is_zero() {
129                        return Err(CodegenError { kind: CodegenErrorKind::DivisionByZero, span: span.clone_box(), token: None });
130                    }
131                    l / r
132                }
133                BinaryOp::Mod => {
134                    if r.is_zero() {
135                        return Err(CodegenError { kind: CodegenErrorKind::DivisionByZero, span: span.clone_box(), token: None });
136                    }
137                    l % r
138                }
139                BinaryOp::Eq => {
140                    if l == r {
141                        U256::from(1u64)
142                    } else {
143                        U256::ZERO
144                    }
145                }
146                BinaryOp::NotEq => {
147                    if l != r {
148                        U256::from(1u64)
149                    } else {
150                        U256::ZERO
151                    }
152                }
153                BinaryOp::Lt => {
154                    if l < r {
155                        U256::from(1u64)
156                    } else {
157                        U256::ZERO
158                    }
159                }
160                BinaryOp::Gt => {
161                    if l > r {
162                        U256::from(1u64)
163                    } else {
164                        U256::ZERO
165                    }
166                }
167                BinaryOp::LtEq => {
168                    if l <= r {
169                        U256::from(1u64)
170                    } else {
171                        U256::ZERO
172                    }
173                }
174                BinaryOp::GtEq => {
175                    if l >= r {
176                        U256::from(1u64)
177                    } else {
178                        U256::ZERO
179                    }
180                }
181            };
182
183            Ok(result.to_be_bytes())
184        }
185
186        Expression::Unary { op, expr: inner_expr, .. } => {
187            // Evaluate the operand using this function to handle nested ArgCalls
188            let val = evaluate_expression_with_context(inner_expr, contract, scope_mgr)?;
189            let v = U256::from_be_bytes(val);
190
191            // Perform unary operation
192            let result = match op {
193                UnaryOp::Neg => U256::ZERO.wrapping_sub(v),
194                UnaryOp::Not => {
195                    if v.is_zero() {
196                        U256::from(1u64)
197                    } else {
198                        U256::ZERO
199                    }
200                }
201            };
202
203            Ok(result.to_be_bytes())
204        }
205
206        Expression::Grouped { expr: inner_expr, .. } => {
207            // Just evaluate the inner expression
208            evaluate_expression_with_context(inner_expr, contract, scope_mgr)
209        }
210
211        // For Literal, Constant, and other types without ArgCalls, use standard evaluation
212        _ => contract.evaluate_constant_expression(expr),
213    }
214}
215
216/// Generates the respective Bytecode for a given Statement
217/// Returns a tuple of (bytes, spans) where spans are optional source positions
218#[allow(clippy::too_many_arguments)]
219pub fn statement_gen<'a>(
220    evm_version: &EVMVersion,
221    s: &Statement,
222    contract: &'a Contract,
223    macro_def: &'a MacroDefinition,
224    scope_mgr: &mut ScopeManager<'a>,
225    offset: &mut usize,
226    jump_table: &mut JumpTable,
227    label_indices: &mut LabelIndices,
228    table_instances: &mut Jumps,
229    utilized_tables: &mut Vec<TableDefinition>,
230    embedded_tables: &mut BTreeMap<String, usize>,
231    starting_offset: usize,
232    relax_jumps: bool,
233) -> Result<StatementGenResult, CodegenError> {
234    let mut bytes = BytecodeSegments::new();
235    let mut spans = vec![];
236
237    tracing::debug!(target: "codegen", "Got Statement: {}", s.ty);
238
239    match &s.ty {
240        StatementType::MacroInvocation(mi) => {
241            // Get the macro definition that matches the name of this invocation
242            let ir_macro = if let Some(m) = contract.find_macro_by_name(&mi.macro_name) {
243                m
244            } else {
245                tracing::error!(
246                    target: "codegen",
247                    "MISSING MACRO INVOCATION \"{}\"",
248                    mi.macro_name
249                );
250                return Err(CodegenError {
251                    kind: CodegenErrorKind::InvalidMacroInvocation(mi.macro_name.clone()),
252                    span: mi.span.clone_box(),
253                    token: None,
254                });
255            };
256
257            tracing::info!(target: "codegen", "FOUND INNER MACRO: {}", ir_macro.name);
258
259            if ir_macro.parameters.len() != mi.args.len() {
260                tracing::error!(target: "codegen", "MISMATCHED ARGUMENT COUNT FOR {}: {} != {}", ir_macro.name, ir_macro.parameters.len(), mi.args.len());
261                return Err(CodegenError {
262                    kind: CodegenErrorKind::InvalidArgumentCount(ir_macro.name.clone(), ir_macro.parameters.len(), mi.args.len()),
263                    span: mi.span.clone_box(),
264                    token: None,
265                });
266            }
267
268            // Tests may not be invoked
269            if ir_macro.test {
270                tracing::error!(target: "codegen", "Tests may not be invoked: {}", ir_macro.name);
271                return Err(CodegenError {
272                    kind: CodegenErrorKind::TestInvocation(ir_macro.name.clone()),
273                    span: ir_macro.span.clone_box(),
274                    token: None,
275                });
276            }
277
278            // Check for circular recursion - if this macro is already in the scope, it's recursive
279            // This must be checked for both outlined and inlined macros
280            if scope_mgr.macro_stack().iter().any(|m| m.name == ir_macro.name) {
281                tracing::error!(
282                    target: "codegen",
283                    "CIRCULAR RECURSION DETECTED: Macro \"{}\" is already in the call stack",
284                    ir_macro.name
285                );
286                return Err(CodegenError {
287                    kind: CodegenErrorKind::CircularMacroInvocation(ir_macro.name.clone()),
288                    span: mi.span.clone_box(),
289                    token: None,
290                });
291            }
292
293            // Check recursion depth limit to prevent stack overflow
294            const MAX_RECURSION_DEPTH: usize = 100;
295            if scope_mgr.depth() >= MAX_RECURSION_DEPTH {
296                tracing::error!(
297                    target: "codegen",
298                    "Maximum recursion depth of {} exceeded",
299                    MAX_RECURSION_DEPTH
300                );
301                return Err(CodegenError {
302                    kind: CodegenErrorKind::CircularMacroInvocation(ir_macro.name.clone()),
303                    span: mi.span.clone_box(),
304                    token: None,
305                });
306            }
307
308            // If invoked macro is a function (outlined), insert a jump to the function's code and a
309            // jumpdest to return to. If it is inlined, insert the macro's code at the
310            // current offset.
311            if ir_macro.outlined {
312                // Get necessary swap ops to reorder stack
313                // PC of the return jumpdest should be below the function's stack inputs
314                let stack_swaps = (0..ir_macro.takes).rev().map(|i| format!("{:02x}", 0x90 + i)).collect::<Vec<_>>();
315
316                // Insert a jump to the outlined macro's code
317                // Functions (outlined macros) are defined at the top-level scope, so use that scope
318                let scope_id = huff_neo_utils::scope::ScopeId::new(vec![], 0);
319                jump_table.insert(
320                    *offset + stack_swaps.len() + 3, // PUSH2 + 2 bytes + stack_swaps.len()
321                    vec![Jump { label: format!("goto_{}", ir_macro.name), bytecode_index: 0, span: s.span.clone(), scope_id }],
322                );
323
324                // Get span info for this statement
325                // Convert from original file coordinates to flattened coordinates
326                let span_info = if !s.span.0.is_empty() {
327                    s.span.0.first().and_then(|sp| {
328                        if sp.start != 0 || sp.end != 0 {
329                            // Convert original span to flattened coordinates
330                            contract.map_original_span_to_flattened(sp)
331                        } else {
332                            None
333                        }
334                    })
335                } else {
336                    None
337                };
338
339                // Store return JUMPDEST PC on the stack and re-order the stack so that
340                // the return JUMPDEST PC is below the function's stack inputs
341                bytes.push_with_offset(
342                    *offset,
343                    Bytes::Raw(format!("{}{:04x}{}", Opcode::Push2, *offset + stack_swaps.len() + 7, stack_swaps.join(""))),
344                );
345                spans.push(span_info);
346                // Insert jump to outlined macro + jumpdest to return to
347                bytes.push_with_offset(
348                    *offset + stack_swaps.len() + 3, // PUSH2 + 2 bytes + stack_swaps.len()
349                    Bytes::JumpPlaceholder(JumpPlaceholderData::new(
350                        format!("goto_{}", ir_macro.name),
351                        PushOpcode::Push2,
352                        format!("{}{}", Opcode::Jump, Opcode::Jumpdest),
353                    )),
354                );
355                spans.push(span_info);
356                // PUSH2 + 2 bytes + stack_swaps.len() + PUSH2 + 2 bytes + JUMP + JUMPDEST
357                *offset += stack_swaps.len() + 8;
358            } else {
359                // Resolve any ArgCall and ArgCallMacroInvocation arguments BEFORE pushing to scope
360                // This ensures we check the calling macro's parameters, not the invoked macro's
361                let mut resolved_args = Vec::new();
362                for arg in &mi.args {
363                    match arg {
364                        MacroArg::ArgCall(arg_call) => {
365                            // Resolve this ArgCall by looking it up in the invocation stack
366                            // An ArgCall in macro invocation arguments (like M2(<arg>)) must be accessible
367                            // from the CURRENT macro's scope - either as a parameter or through invocation chain
368                            let mut resolved = None;
369
370                            // First check if current macro has this parameter
371                            let current_has_param = macro_def.parameters.iter().any(|p| p.name.as_deref() == Some(&arg_call.name));
372
373                            // If current macro doesn't have this parameter, check if it was passed to current macro
374                            if !current_has_param {
375                                // The ArgCall must have been explicitly passed to the current macro
376                                // Check if the current macro's invocation (last in mis) contains this ArgCall
377                                let mut passed_to_current = false;
378
379                                if let Some(current_mi) = scope_mgr.current_invocation() {
380                                    // Check if any of the arguments passed to current macro is this ArgCall
381                                    fn contains_arg_call(args: &[MacroArg], name: &str) -> bool {
382                                        args.iter().any(|arg| match arg {
383                                            MacroArg::ArgCall(ac) => ac.name == name,
384                                            MacroArg::MacroCall(mc) => contains_arg_call(&mc.args, name),
385                                            _ => false,
386                                        })
387                                    }
388
389                                    passed_to_current = contains_arg_call(&current_mi.args, &arg_call.name);
390                                }
391
392                                // If not passed to current macro, error
393                                if !passed_to_current {
394                                    return Err(CodegenError {
395                                        kind: CodegenErrorKind::MissingArgumentDefinition(arg_call.name.clone()),
396                                        span: mi.span.clone_box(),
397                                        token: None,
398                                    });
399                                }
400                            }
401
402                            // Search through the invocation stack for the argument's value
403                            let mis = scope_mgr.invocation_vec();
404                            for (_, parent_mi) in mis.iter().rev() {
405                                if let Some(parent_macro) = contract.find_macro_by_name(&parent_mi.macro_name)
406                                    && let Some(param_idx) =
407                                        parent_macro.parameters.iter().position(|p| p.name.as_deref() == Some(&arg_call.name))
408                                    && let Some(arg_value) = parent_mi.args.get(param_idx)
409                                {
410                                    resolved = Some(arg_value.clone());
411                                    break;
412                                }
413                            }
414
415                            // Use resolved value or keep the ArgCall if not found
416                            resolved_args.push(resolved.unwrap_or_else(|| arg.clone()));
417                        }
418                        MacroArg::ArgCallMacroInvocation(arg_name, invoc_args) => {
419                            // Resolve the argument name to find the actual macro
420                            let mis = scope_mgr.invocation_vec();
421                            let mut resolved_macro_name = None;
422
423                            // Search through the invocation stack for the argument's value
424                            for (_, parent_mi) in mis.iter().rev() {
425                                if let Some(parent_macro) = contract.find_macro_by_name(&parent_mi.macro_name)
426                                    && let Some(param_idx) =
427                                        parent_macro.parameters.iter().position(|p| p.name.as_deref() == Some(arg_name))
428                                {
429                                    // Check what type of argument value we have
430                                    if let Some(arg_value) = parent_mi.args.get(param_idx) {
431                                        match arg_value {
432                                            MacroArg::Ident(macro_name) => {
433                                                // Direct macro name reference
434                                                resolved_macro_name = Some(macro_name.clone());
435                                                break;
436                                            }
437                                            MacroArg::MacroCall(macro_call) => {
438                                                // The argument is a macro invocation - extract the macro name
439                                                resolved_macro_name = Some(macro_call.macro_name.clone());
440                                                break;
441                                            }
442                                            _ => {
443                                                // Other types not supported for macro invocation, continue searching
444                                                continue;
445                                            }
446                                        }
447                                    }
448                                }
449                            }
450
451                            // If we resolved the macro name, create a regular MacroCall
452                            if let Some(macro_name) = resolved_macro_name {
453                                resolved_args.push(MacroArg::MacroCall(MacroInvocation {
454                                    macro_name,
455                                    args: invoc_args.clone(),
456                                    span: AstSpan(vec![]),
457                                }));
458                            } else {
459                                // Keep the original if we couldn't resolve
460                                resolved_args.push(arg.clone());
461                            }
462                        }
463                        MacroArg::MacroCall(macro_call) => {
464                            // Recursively resolve any ArgCalls in the macro call's arguments
465                            fn resolve_macro_call_args<'a>(
466                                macro_call: &MacroInvocation,
467                                scope_mgr: &ScopeManager<'a>,
468                                contract: &Contract,
469                                macro_def: &MacroDefinition,
470                            ) -> Result<Vec<MacroArg>, CodegenError> {
471                                let mis = scope_mgr.invocation_vec();
472                                let mut resolved_args = Vec::new();
473                                for arg in &macro_call.args {
474                                    match arg {
475                                        MacroArg::ArgCall(arg_call) => {
476                                            // First check if current macro has this parameter
477                                            let current_has_param =
478                                                macro_def.parameters.iter().any(|p| p.name.as_deref() == Some(&arg_call.name));
479
480                                            // If current macro doesn't have this parameter, check if it was passed to current macro
481                                            if !current_has_param {
482                                                // The ArgCall must have been explicitly passed to the current macro
483                                                // Check if the current macro's invocation (last in mis) contains this ArgCall
484                                                let mut passed_to_current = false;
485
486                                                if let Some((_, current_mi)) = mis.last() {
487                                                    // Check if any of the arguments passed to current macro is this ArgCall
488                                                    fn contains_arg_call(args: &[MacroArg], name: &str) -> bool {
489                                                        args.iter().any(|arg| match arg {
490                                                            MacroArg::ArgCall(ac) => ac.name == name,
491                                                            MacroArg::MacroCall(mc) => contains_arg_call(&mc.args, name),
492                                                            _ => false,
493                                                        })
494                                                    }
495
496                                                    passed_to_current = contains_arg_call(&current_mi.args, &arg_call.name);
497                                                }
498
499                                                // If not passed to current macro, error
500                                                if !passed_to_current {
501                                                    return Err(CodegenError {
502                                                        kind: CodegenErrorKind::MissingArgumentDefinition(arg_call.name.clone()),
503                                                        span: macro_call.span.clone_box(),
504                                                        token: None,
505                                                    });
506                                                }
507                                            }
508
509                                            // Resolve this ArgCall by looking it up in the current context
510                                            let mut resolved = None;
511                                            for (_, parent_mi) in mis.iter().rev() {
512                                                if let Some(parent_macro) = contract.find_macro_by_name(&parent_mi.macro_name)
513                                                    && let Some(param_idx) = parent_macro
514                                                        .parameters
515                                                        .iter()
516                                                        .position(|p| p.name.as_deref() == Some(&arg_call.name))
517                                                    && let Some(arg_value) = parent_mi.args.get(param_idx)
518                                                {
519                                                    resolved = Some(arg_value.clone());
520                                                    break;
521                                                }
522                                            }
523                                            resolved_args.push(resolved.unwrap_or_else(|| arg.clone()));
524                                        }
525                                        MacroArg::MacroCall(nested_macro_call) => {
526                                            // Recursively resolve nested MacroCalls
527                                            let nested_resolved_args =
528                                                resolve_macro_call_args(nested_macro_call, scope_mgr, contract, macro_def)?;
529                                            resolved_args.push(MacroArg::MacroCall(MacroInvocation {
530                                                macro_name: nested_macro_call.macro_name.clone(),
531                                                args: nested_resolved_args,
532                                                span: nested_macro_call.span.clone(),
533                                            }));
534                                        }
535                                        _ => resolved_args.push(arg.clone()),
536                                    }
537                                }
538                                Ok(resolved_args)
539                            }
540
541                            let resolved_macro_args = resolve_macro_call_args(macro_call, scope_mgr, contract, macro_def)?;
542                            resolved_args.push(MacroArg::MacroCall(MacroInvocation {
543                                macro_name: macro_call.macro_name.clone(),
544                                args: resolved_macro_args,
545                                span: macro_call.span.clone(),
546                            }));
547                        }
548                        _ => resolved_args.push(arg.clone()),
549                    }
550                }
551
552                // Create a resolved invocation
553                let resolved_mi = MacroInvocation { macro_name: mi.macro_name.clone(), args: resolved_args, span: mi.span.clone() };
554
555                // Push the invoked macro to the scope AFTER resolving arguments
556                // This is important for circular recursion detection
557                scope_mgr.push_macro_with_invocation(ir_macro, *offset, resolved_mi);
558
559                let mut res: BytecodeRes =
560                    match Codegen::macro_to_bytecode(evm_version, ir_macro, contract, scope_mgr, *offset, relax_jumps) {
561                        Ok(r) => r,
562                        Err(e) => {
563                            tracing::error!(
564                                target: "codegen",
565                                "FAILED TO RECURSE INTO MACRO \"{}\"",
566                                ir_macro.name
567                            );
568                            scope_mgr.pop_macro();
569                            return Err(e);
570                        }
571                    };
572
573                // Note: The macro is already popped by macro_to_bytecode at the end
574                // because depth > 1 when processing inline macros recursively. We should NOT pop again here.
575
576                // Set jump table values
577                tracing::debug!(target: "codegen", "Unmatched jumps: {:?}", res.unmatched_jumps.iter().map(|uj| uj.label.clone()).collect::<Vec<String>>());
578                for j in res.unmatched_jumps.iter_mut() {
579                    let new_index = j.bytecode_index;
580                    j.bytecode_index = 0;
581                    let mut new_jumps = if let Some(jumps) = jump_table.get(&new_index) { jumps.clone() } else { vec![] };
582                    new_jumps.push(j.clone());
583                    jump_table.insert(new_index, new_jumps);
584                }
585                tracing::debug!(target: "codegen", "Extending table_instances with {} items from macro {}", res.table_instances.len(), ir_macro.name);
586                table_instances.extend(res.table_instances);
587                label_indices.extend(res.label_indices);
588                // Only add tables that are not already in the utilized tables
589                let res_unique_tables =
590                    res.utilized_tables.iter().filter(|t| !utilized_tables.contains(t)).cloned().collect::<Vec<TableDefinition>>();
591                utilized_tables.extend(res_unique_tables);
592
593                // Increase offset by byte length of recursed macro
594                *offset += res.bytes.iter().map(|seg| seg.bytes.len()).sum::<usize>();
595                // Add the macro's bytecode and spans to the final result
596                bytes.extend(res.bytes);
597                // Preserve the spans from the nested macro expansion
598                spans.extend(res.spans)
599            }
600        }
601        StatementType::ImmediateOpcode { opcode, operands } => {
602            // Resolve operands in the current macro invocation context (constants, args, loop values)
603            let mut values = Vec::with_capacity(operands.len());
604            for operand in operands {
605                values.push(U256::from_be_bytes(evaluate_expression_with_context(operand, contract, scope_mgr)?));
606            }
607            let immediate = opcode.encode_immediate(&values).map_err(|msg| CodegenError {
608                kind: CodegenErrorKind::InvalidOpcodeOperand(msg),
609                span: s.span.clone_box(),
610                token: None,
611            })?;
612            let code = format!("{opcode}{immediate}");
613            tracing::info!(target: "codegen", "IMMEDIATE OPCODE: {:?} {:?} => {}", opcode, values, code);
614
615            bytes.push_with_offset(*offset, Bytes::Raw(code.clone()));
616            let span_info = s
617                .span
618                .0
619                .first()
620                .and_then(|sp| if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None });
621            spans.push(span_info);
622            *offset += code.len() / 2;
623        }
624        StatementType::Label(label) => {
625            // Add JUMPDEST opcode to final result and add to label_indices
626            tracing::info!(target: "codegen", "RECURSE BYTECODE GOT LABEL: {:?}", label.name);
627
628            // Build the scope path from the current macro invocation stack
629            // Use offset to make each invocation unique
630            let scope_id = scope_mgr.current_scope();
631
632            // Try to insert the label with scope information
633            if let Err(err_msg) = label_indices.insert(label.name.clone(), *offset, &scope_id) {
634                tracing::error!(target: "codegen", "DUPLICATE LABEL: {}", err_msg);
635                return Err(CodegenError {
636                    kind: CodegenErrorKind::DuplicateLabelInScope(label.name.clone()),
637                    span: s.span.clone_box(),
638                    token: None,
639                });
640            }
641
642            bytes.push_with_offset(*offset, Bytes::Raw(Opcode::Jumpdest.to_string()));
643            // Add span for the JUMPDEST
644            let span_info = if !s.span.0.is_empty() {
645                s.span.0.first().and_then(|sp| {
646                    if sp.start != 0 || sp.end != 0 {
647                        // Convert original span to flattened coordinates
648                        contract.map_original_span_to_flattened(sp)
649                    } else {
650                        None
651                    }
652                })
653            } else {
654                None
655            };
656            spans.push(span_info);
657            *offset += 1;
658        }
659        StatementType::LabelCall(label) => {
660            // Generate code for a `LabelCall`
661            // PUSH2 + 2 byte destination (placeholder for now, filled in `Codegen::fill_unmatched`
662            tracing::info!(target: "codegen", "RECURSE BYTECODE GOT LABEL CALL: {}", label);
663
664            // Get the scope for this jump
665            let scope_id = scope_mgr.current_scope();
666
667            jump_table.insert(*offset, vec![Jump { label: label.to_string(), bytecode_index: 0, span: s.span.clone(), scope_id }]);
668            bytes.push_with_offset(
669                *offset,
670                Bytes::JumpPlaceholder(JumpPlaceholderData::new(label.to_string(), PushOpcode::Push2, String::new())),
671            );
672            // Add span for the PUSH2
673            let span_info = if !s.span.0.is_empty() {
674                s.span.0.first().and_then(|sp| {
675                    if sp.start != 0 || sp.end != 0 {
676                        // Convert original span to flattened coordinates
677                        contract.map_original_span_to_flattened(sp)
678                    } else {
679                        None
680                    }
681                })
682            } else {
683                None
684            };
685            spans.push(span_info);
686            *offset += 3;
687        }
688        StatementType::BuiltinFunctionCall(bf) => {
689            let bytes_before = bytes.len();
690            builtin_function::builtin_function_gen(
691                evm_version,
692                contract,
693                macro_def,
694                scope_mgr,
695                offset,
696                table_instances,
697                utilized_tables,
698                embedded_tables,
699                starting_offset,
700                &mut bytes,
701                bf,
702                relax_jumps,
703            )?;
704            // Add spans for the bytes added by builtin function
705            let bytes_added = bytes.len() - bytes_before;
706            let span_info = if !s.span.0.is_empty() {
707                s.span.0.first().and_then(|sp| {
708                    if sp.start != 0 || sp.end != 0 {
709                        // Convert original span to flattened coordinates
710                        contract.map_original_span_to_flattened(sp)
711                    } else {
712                        None
713                    }
714                })
715            } else {
716                None
717            };
718            for _ in 0..bytes_added {
719                spans.push(span_info);
720            }
721        }
722        StatementType::ArgMacroInvocation(_parent_macro_name, arg_name, args) => {
723            // Handle macro invocation through argument
724            // We need to resolve the argument to get the actual macro name
725            tracing::info!(target: "codegen", "Processing ArgMacroInvocation: <{}>()", arg_name);
726
727            // Function to recursively resolve an argument through the invocation stack
728            fn resolve_argument<'a>(arg_name: &str, scope_mgr: &ScopeManager<'a>, contract: &Contract) -> Result<String, CodegenErrorKind> {
729                let mis = scope_mgr.invocation_vec();
730                // Search through the macro invocation stack
731                for (_, mi) in mis.iter().rev() {
732                    // Find the macro definition
733                    if let Some(invoked_macro) = contract.find_macro_by_name(&mi.macro_name) {
734                        // Check if this macro has our argument
735                        if let Some(param_index) = invoked_macro.parameters.iter().position(|p| p.name.as_deref() == Some(arg_name)) {
736                            // Get the corresponding argument value
737                            if let Some(arg_value) = mi.args.get(param_index) {
738                                match arg_value {
739                                    MacroArg::Ident(macro_name) => {
740                                        // Direct macro name - we're done
741                                        return Ok(macro_name.clone());
742                                    }
743                                    MacroArg::ArgCall(arg_call) => {
744                                        // The argument is itself an argument reference - need to resolve it recursively
745                                        return resolve_argument(&arg_call.name, scope_mgr, contract);
746                                    }
747                                    MacroArg::ArgCallMacroInvocation(inner_arg_name, _) => {
748                                        // The argument is an invocation of another argument - resolve that argument
749                                        return resolve_argument(inner_arg_name, scope_mgr, contract);
750                                    }
751                                    MacroArg::Noop => {
752                                        // __NOOP cannot be invoked as a macro
753                                        return Err(CodegenErrorKind::InvalidMacroArgumentType(format!(
754                                            "Cannot invoke __NOOP as a macro in argument '{}'",
755                                            arg_name
756                                        )));
757                                    }
758                                    _ => {
759                                        // Other types of arguments not supported for macro invocation
760                                        continue;
761                                    }
762                                }
763                            }
764                        }
765                    }
766                }
767                Err(CodegenErrorKind::MissingArgumentDefinition(arg_name.to_string()))
768            }
769
770            // Resolve the argument to get the actual macro name
771            let macro_name = match resolve_argument(arg_name, scope_mgr, contract) {
772                Ok(name) => {
773                    tracing::info!(target: "codegen", "Resolved <{}> to macro: {}", arg_name, name);
774                    name
775                }
776                Err(error_kind) => {
777                    tracing::error!(target: "codegen", "Failed to resolve argument '{}': {:?}", arg_name, error_kind);
778                    return Err(CodegenError { kind: error_kind, span: s.span.clone_box(), token: None });
779                }
780            };
781
782            // Create a MacroInvocation with the resolved macro name
783            let resolved_invocation = MacroInvocation { macro_name, args: args.clone(), span: s.span.clone() };
784
785            // Process it as a regular macro invocation
786            let macro_invocation_statement =
787                Statement { ty: StatementType::MacroInvocation(resolved_invocation.clone()), span: s.span.clone() };
788
789            // Recursively process the resolved macro invocation
790            return statement_gen(
791                evm_version,
792                &macro_invocation_statement,
793                contract,
794                macro_def,
795                scope_mgr,
796                offset,
797                jump_table,
798                label_indices,
799                table_instances,
800                utilized_tables,
801                embedded_tables,
802                starting_offset,
803                relax_jumps,
804            );
805        }
806        StatementType::IfStatement { condition, then_branch, else_if_branches, else_branch } => {
807            // Evaluate condition with macro invocation context (supports <arg>)
808            let condition_val = evaluate_expression_with_context(condition, contract, scope_mgr)?;
809            let condition_true = !U256::from_be_bytes(condition_val).is_zero();
810
811            // Select the appropriate branch
812            let selected_branch = if condition_true {
813                then_branch.as_slice()
814            } else {
815                // Check else if branches
816                let mut matched_branch = None;
817                for (else_if_condition, else_if_body) in else_if_branches {
818                    let else_if_val = evaluate_expression_with_context(else_if_condition, contract, scope_mgr)?;
819                    let else_if_true = !U256::from_be_bytes(else_if_val).is_zero();
820
821                    if else_if_true {
822                        matched_branch = Some(else_if_body.as_slice());
823                        break;
824                    }
825                }
826
827                // Use matched else if branch, or else branch, or empty
828                matched_branch.or(else_branch.as_ref().map(|b| b.as_slice())).unwrap_or(&[])
829            };
830
831            // Convert branch statements to IRBytes first (like labels do)
832            let branch_irbytes = MacroDefinition::to_irbytes(evm_version, selected_branch);
833
834            // Process each IRByte through statement_gen or directly as bytes
835            for ir_byte in branch_irbytes {
836                match &ir_byte.ty {
837                    IRByteType::Statement(stmt) => {
838                        let (stmt_bytes, mut stmt_spans) = statement_gen(
839                            evm_version,
840                            stmt,
841                            contract,
842                            macro_def,
843                            scope_mgr,
844                            offset,
845                            jump_table,
846                            label_indices,
847                            table_instances,
848                            utilized_tables,
849                            embedded_tables,
850                            starting_offset,
851                            relax_jumps,
852                        )?;
853                        bytes.extend(stmt_bytes);
854                        spans.append(&mut stmt_spans);
855                    }
856                    IRByteType::Bytes(b) => {
857                        *offset += b.len();
858                        bytes.push_with_offset(starting_offset, b.to_owned());
859                        let span_info = if !ir_byte.span.0.is_empty() {
860                            ir_byte.span.0.first().and_then(|sp| {
861                                if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None }
862                            })
863                        } else {
864                            None
865                        };
866                        spans.push(span_info);
867                    }
868                    IRByteType::Constant(name) => {
869                        // Handle constant references
870                        if let Some(push_value) = constant_gen(name, contract, ir_byte.span)? {
871                            let push_bytes = push_value.to_hex_with_opcode(evm_version);
872                            *offset += push_bytes.len() / 2;
873                            bytes.push_with_offset(starting_offset, Bytes::Raw(push_bytes));
874                            let span_info = if !ir_byte.span.0.is_empty() {
875                                ir_byte.span.0.first().and_then(|sp| {
876                                    if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None }
877                                })
878                            } else {
879                                None
880                            };
881                            spans.push(span_info);
882                        }
883                    }
884                    IRByteType::ArgCall(parent_macro_name, arg_name) => {
885                        // Handle arg calls
886                        let bytes_before = bytes.len();
887                        bubble_arg_call(
888                            evm_version,
889                            parent_macro_name.to_owned(),
890                            arg_name.to_owned(),
891                            &mut bytes,
892                            macro_def,
893                            contract,
894                            scope_mgr,
895                            offset,
896                            jump_table,
897                            table_instances,
898                            utilized_tables,
899                            ir_byte.span,
900                            relax_jumps,
901                        )?;
902                        let bytes_added = bytes.len() - bytes_before;
903                        let span_info = if !ir_byte.span.0.is_empty() {
904                            ir_byte.span.0.first().and_then(|sp| {
905                                if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None }
906                            })
907                        } else {
908                            None
909                        };
910                        for _ in 0..bytes_added {
911                            spans.push(span_info);
912                        }
913                    }
914                }
915            }
916        }
917        StatementType::ForLoop { variable, start, end, step, body } => {
918            // Evaluate loop bounds with macro invocation context (supports <arg>)
919            let start_val = evaluate_expression_with_context(start, contract, scope_mgr)?;
920            let end_val = evaluate_expression_with_context(end, contract, scope_mgr)?;
921            let step_val = if let Some(step_expr) = step {
922                evaluate_expression_with_context(step_expr, contract, scope_mgr)?
923            } else {
924                str_to_bytes32("1")
925            };
926
927            // Convert to U256 for iteration
928            let start_u256 = U256::from_be_bytes(start_val);
929            let end_u256 = U256::from_be_bytes(end_val);
930            let step_u256 = U256::from_be_bytes(step_val);
931
932            if step_u256.is_zero() {
933                return Err(CodegenError { kind: CodegenErrorKind::InvalidLoopStep, span: s.span.clone_box(), token: None });
934            }
935
936            // Maximum iterations safety check
937            const MAX_LOOP_ITERATIONS: usize = 100_000;
938            let mut current = start_u256;
939            let mut iteration_count = 0;
940
941            // Unroll the loop at compile-time
942            while current < end_u256 {
943                if iteration_count >= MAX_LOOP_ITERATIONS {
944                    return Err(CodegenError {
945                        kind: CodegenErrorKind::LoopIterationLimitExceeded(MAX_LOOP_ITERATIONS),
946                        span: s.span.clone_box(),
947                        token: None,
948                    });
949                }
950
951                // Substitute loop variable in the body
952                let current_bytes = current.to_be_bytes::<32>();
953                let substituted_body = substitute_loop_variable_in_statements(body, variable, &current_bytes)?;
954
955                // Convert substituted body to IRBytes first (like labels and if statements)
956                let loop_irbytes = MacroDefinition::to_irbytes(evm_version, &substituted_body);
957
958                // Process each IRByte
959                for ir_byte in loop_irbytes {
960                    match &ir_byte.ty {
961                        IRByteType::Statement(stmt) => {
962                            let (stmt_bytes, mut stmt_spans) = statement_gen(
963                                evm_version,
964                                stmt,
965                                contract,
966                                macro_def,
967                                scope_mgr,
968                                offset,
969                                jump_table,
970                                label_indices,
971                                table_instances,
972                                utilized_tables,
973                                embedded_tables,
974                                starting_offset,
975                                relax_jumps,
976                            )?;
977                            bytes.extend(stmt_bytes);
978                            spans.append(&mut stmt_spans);
979                        }
980                        IRByteType::Bytes(b) => {
981                            *offset += b.len();
982                            bytes.push_with_offset(starting_offset, b.to_owned());
983                            let span_info = if !ir_byte.span.0.is_empty() {
984                                ir_byte.span.0.first().and_then(|sp| {
985                                    if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None }
986                                })
987                            } else {
988                                None
989                            };
990                            spans.push(span_info);
991                        }
992                        IRByteType::Constant(name) => {
993                            if let Some(push_value) = constant_gen(name, contract, ir_byte.span)? {
994                                let push_bytes = push_value.to_hex_with_opcode(evm_version);
995                                *offset += push_bytes.len() / 2;
996                                bytes.push_with_offset(starting_offset, Bytes::Raw(push_bytes));
997                                let span_info = if !ir_byte.span.0.is_empty() {
998                                    ir_byte.span.0.first().and_then(|sp| {
999                                        if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None }
1000                                    })
1001                                } else {
1002                                    None
1003                                };
1004                                spans.push(span_info);
1005                            }
1006                        }
1007                        IRByteType::ArgCall(parent_macro_name, arg_name) => {
1008                            let bytes_before = bytes.len();
1009                            bubble_arg_call(
1010                                evm_version,
1011                                parent_macro_name.to_owned(),
1012                                arg_name.to_owned(),
1013                                &mut bytes,
1014                                macro_def,
1015                                contract,
1016                                scope_mgr,
1017                                offset,
1018                                jump_table,
1019                                table_instances,
1020                                utilized_tables,
1021                                ir_byte.span,
1022                                relax_jumps,
1023                            )?;
1024                            let bytes_added = bytes.len() - bytes_before;
1025                            let span_info = if !ir_byte.span.0.is_empty() {
1026                                ir_byte.span.0.first().and_then(|sp| {
1027                                    if sp.start != 0 || sp.end != 0 { contract.map_original_span_to_flattened(sp) } else { None }
1028                                })
1029                            } else {
1030                                None
1031                            };
1032                            for _ in 0..bytes_added {
1033                                spans.push(span_info);
1034                            }
1035                        }
1036                    }
1037                }
1038
1039                current = current.saturating_add(step_u256);
1040                iteration_count += 1;
1041            }
1042        }
1043        sty => {
1044            tracing::error!(target: "codegen", "CURRENT MACRO DEF: {}", macro_def.name);
1045            tracing::error!(target: "codegen", "UNEXPECTED STATEMENT: {:?}", sty);
1046            return Err(CodegenError { kind: CodegenErrorKind::InvalidMacroStatement, span: s.span.clone_box(), token: None });
1047        }
1048    }
1049
1050    Ok((bytes, spans))
1051}
1052
1053/// Substitute loop variable in a list of statements
1054fn substitute_loop_variable_in_statements(
1055    statements: &[Statement],
1056    var_name: &str,
1057    value: &[u8; 32],
1058) -> Result<Vec<Statement>, CodegenError> {
1059    let mut substituted = Vec::new();
1060
1061    for statement in statements {
1062        let new_statement = substitute_loop_variable_in_statement(statement, var_name, value)?;
1063        substituted.push(new_statement);
1064    }
1065
1066    Ok(substituted)
1067}
1068
1069/// Substitute loop variable in a single statement
1070fn substitute_loop_variable_in_statement(statement: &Statement, var_name: &str, value: &[u8; 32]) -> Result<Statement, CodegenError> {
1071    let new_ty = match &statement.ty {
1072        StatementType::Constant(name) => {
1073            // Check if this is a loop variable placeholder
1074            if name == &format!("__LOOP_VAR_{}", var_name) {
1075                // Replace with literal
1076                StatementType::Literal(*value)
1077            } else {
1078                statement.ty.clone()
1079            }
1080        }
1081        StatementType::Label(label) => {
1082            // Recursively substitute in label body and rename label to prevent collisions
1083            let substituted_inner = substitute_loop_variable_in_statements(&label.inner, var_name, value)?;
1084            let iter_value = U256::from_be_bytes(*value);
1085            let mut new_label = label.clone();
1086            new_label.inner = substituted_inner;
1087            new_label.name = format!("{}_{}", label.name, iter_value);
1088            StatementType::Label(new_label)
1089        }
1090        StatementType::ImmediateOpcode { opcode, operands } => StatementType::ImmediateOpcode {
1091            opcode: *opcode,
1092            operands: operands.iter().map(|o| substitute_in_expression(o, var_name, value)).collect(),
1093        },
1094        StatementType::BuiltinFunctionCall(bf) => StatementType::BuiltinFunctionCall(BuiltinFunctionCall {
1095            args: bf
1096                .args
1097                .iter()
1098                .map(|arg| match arg {
1099                    BuiltinFunctionArg::Expression(expr) => BuiltinFunctionArg::Expression(substitute_in_expression(expr, var_name, value)),
1100                    arg => arg.clone(),
1101                })
1102                .collect(),
1103            ..bf.clone()
1104        }),
1105        StatementType::MacroInvocation(mi) => {
1106            StatementType::MacroInvocation(MacroInvocation { args: substitute_in_macro_args(&mi.args, var_name, value), ..mi.clone() })
1107        }
1108        StatementType::ArgMacroInvocation(parent_macro_name, arg_name, args) => {
1109            StatementType::ArgMacroInvocation(parent_macro_name.clone(), arg_name.clone(), substitute_in_macro_args(args, var_name, value))
1110        }
1111        StatementType::LabelCall(label_name) => {
1112            // Update label references to match renamed labels
1113            let iter_value = U256::from_be_bytes(*value);
1114            StatementType::LabelCall(format!("{}_{}", label_name, iter_value))
1115        }
1116        StatementType::ForLoop { variable: loop_var, start, end, step, body } => {
1117            // Recursively substitute in nested for loop body
1118            let substituted_body = substitute_loop_variable_in_statements(body, var_name, value)?;
1119            StatementType::ForLoop {
1120                variable: loop_var.clone(),
1121                start: substitute_in_expression(start, var_name, value),
1122                end: substitute_in_expression(end, var_name, value),
1123                step: step.as_ref().map(|s| substitute_in_expression(s, var_name, value)),
1124                body: substituted_body,
1125            }
1126        }
1127        StatementType::IfStatement { condition, then_branch, else_if_branches, else_branch } => {
1128            // Recursively substitute in if statement branches
1129            let substituted_then = substitute_loop_variable_in_statements(then_branch, var_name, value)?;
1130            let substituted_else_if: Vec<_> = else_if_branches
1131                .iter()
1132                .map(|(cond, body)| {
1133                    Ok((substitute_in_expression(cond, var_name, value), substitute_loop_variable_in_statements(body, var_name, value)?))
1134                })
1135                .collect::<Result<Vec<_>, CodegenError>>()?;
1136            let substituted_else =
1137                else_branch.as_ref().map(|body| substitute_loop_variable_in_statements(body, var_name, value)).transpose()?;
1138
1139            StatementType::IfStatement {
1140                condition: substitute_in_expression(condition, var_name, value),
1141                then_branch: substituted_then,
1142                else_if_branches: substituted_else_if,
1143                else_branch: substituted_else,
1144            }
1145        }
1146        _ => statement.ty.clone(),
1147    };
1148
1149    Ok(Statement { ty: new_ty, span: statement.span.clone() })
1150}
1151
1152/// Substitute loop variable in an expression
1153fn substitute_in_expression(expr: &Expression, var_name: &str, value: &[u8; 32]) -> Expression {
1154    match expr {
1155        Expression::Constant { name, span } => {
1156            // Check if this is a loop variable placeholder constant
1157            if name == &format!("__LOOP_VAR_{}", var_name) {
1158                Expression::Literal { value: *value, span: span.clone() }
1159            } else {
1160                expr.clone()
1161            }
1162        }
1163        Expression::Binary { left, op, right, span } => Expression::Binary {
1164            left: Box::new(substitute_in_expression(left, var_name, value)),
1165            op: op.clone(),
1166            right: Box::new(substitute_in_expression(right, var_name, value)),
1167            span: span.clone(),
1168        },
1169        Expression::Unary { op, expr: inner_expr, span } => {
1170            Expression::Unary { op: op.clone(), expr: Box::new(substitute_in_expression(inner_expr, var_name, value)), span: span.clone() }
1171        }
1172        Expression::Grouped { expr: inner_expr, span } => {
1173            Expression::Grouped { expr: Box::new(substitute_in_expression(inner_expr, var_name, value)), span: span.clone() }
1174        }
1175        _ => expr.clone(),
1176    }
1177}
1178
1179/// Substitute loop variable in macro invocation arguments, including nested invocations
1180fn substitute_in_macro_args(args: &[MacroArg], var_name: &str, value: &[u8; 32]) -> Vec<MacroArg> {
1181    args.iter()
1182        .map(|arg| match arg {
1183            MacroArg::Ident(name) if name == &format!("__LOOP_VAR_{}", var_name) => MacroArg::HexLiteral(*value),
1184            MacroArg::MacroCall(mi) => {
1185                MacroArg::MacroCall(MacroInvocation { args: substitute_in_macro_args(&mi.args, var_name, value), ..mi.clone() })
1186            }
1187            MacroArg::ArgCallMacroInvocation(name, nested) => {
1188                MacroArg::ArgCallMacroInvocation(name.clone(), substitute_in_macro_args(nested, var_name, value))
1189            }
1190            _ => arg.clone(),
1191        })
1192        .collect()
1193}
huff_neo_codegen/irgen/statements.rs

huff_neo_codegen/irgen/
statements.rs
