huff_neo_test_runner/
internal.rs

1use alloy_primitives::{Address, U256};
2use foundry_evm::traces::{CallTraceStep, DecodedInternalCall, DecodedTraceStep, SparsedTraceArena};
3use revm::bytecode::opcode::{JUMP, JUMPDEST};
4
5/// An outlined Huff function (`#define fn`) of a compiled test, which traces show as an internal call
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct HuffFunction {
8    /// The name of the function
9    pub name: String,
10    /// The name of the file that defines the function, which traces show before the name
11    pub file: Option<String>,
12    /// The program counter of the JUMPDEST the function starts at
13    pub entry_pc: usize,
14    /// How many stack items the function takes
15    pub takes: usize,
16    /// How many stack items the function returns
17    pub returns: usize,
18}
19
20/// A call into a Huff function that has not returned yet
21struct OpenCall {
22    /// Index of the function in the list of functions
23    function: usize,
24    /// Index of the step that jumps into the function
25    start: usize,
26    /// Program counter of that jump; the function returns to the JUMPDEST right after it
27    jump_pc: usize,
28}
29
30/// Marks the calls of Huff functions in the traces of the test at `address` as internal calls, so
31/// the trace shows each call with its arguments and return values.
32///
33/// It needs the jump steps that `--decode-internal` records. A call jumps to the function's entry
34/// JUMPDEST and returns to the JUMPDEST right after the calling jump; a call that never returns, for
35/// example because it reverts, is left out.
36pub fn decode_huff_internal_calls(arena: &mut SparsedTraceArena, address: Address, functions: &[HuffFunction]) {
37    if functions.is_empty() {
38        return;
39    }
40
41    for node in arena.arena.nodes_mut().iter_mut().filter(|node| node.trace.address == address && !node.trace.kind.is_any_create()) {
42        let steps = &mut node.trace.steps;
43        let mut open: Vec<OpenCall> = Vec::new();
44
45        for i in 1..steps.len() {
46            if steps[i].op.get() != JUMPDEST || steps[i - 1].op.get() != JUMP {
47                continue;
48            }
49            let pc = steps[i].pc;
50
51            // A jump back to the JUMPDEST after a calling jump returns from that call
52            if let Some(depth) = open.iter().rposition(|call| call.jump_pc + 1 == pc) {
53                let call = &open[depth];
54                let function = &functions[call.function];
55                let decoded = DecodedInternalCall {
56                    func_name: match &function.file {
57                        Some(file) => format!("{file}::{}", function.name),
58                        None => function.name.clone(),
59                    },
60                    args: Some(top_of_stack(&steps[call.start + 1], function.takes)),
61                    return_data: (function.returns > 0).then(|| top_of_stack(&steps[i], function.returns)),
62                };
63                steps[call.start].decoded = Some(Box::new(DecodedTraceStep::InternalCall(decoded, i)));
64                open.truncate(depth);
65                continue;
66            }
67
68            if let Some(function) = functions.iter().position(|function| function.entry_pc == pc) {
69                open.push(OpenCall { function, start: i - 1, jump_pc: steps[i - 1].pc });
70            }
71        }
72    }
73}
74
75/// The top `count` stack items of a step, in the order they were pushed
76fn top_of_stack(step: &CallTraceStep, count: usize) -> Vec<String> {
77    let stack: &[U256] = step.stack.as_deref().unwrap_or_default();
78    stack[stack.len().saturating_sub(count)..].iter().map(|word| format!("{word:#x}")).collect()
79}
huff_neo_test_runner/internal.rs

huff_neo_test_runner/
internal.rs
