foundry_evm_traces/
folded_stack_trace.rs

1use alloy_primitives::hex::ToHexExt;
2use revm_inspectors::tracing::{
3    CallTraceArena,
4    types::{CallTraceNode, CallTraceStep, DecodedTraceStep, TraceMemberOrder},
5};
6
7/// Builds a folded stack trace from a call trace arena.
8pub fn build(arena: &CallTraceArena) -> Vec<String> {
9    let mut fst = EvmFoldedStackTraceBuilder::default();
10    fst.process_call_node(arena.nodes(), 0);
11    fst.build()
12}
13
14/// Wrapper for building a folded stack trace using EVM call trace node.
15#[derive(Default)]
16pub struct EvmFoldedStackTraceBuilder {
17    /// Raw folded stack trace builder.
18    fst: FoldedStackTraceBuilder,
19}
20
21impl EvmFoldedStackTraceBuilder {
22    /// Returns the folded stack trace.
23    pub fn build(self) -> Vec<String> {
24        self.fst.build()
25    }
26
27    /// Creates an entry for a EVM CALL in the folded stack trace. This method recursively processes
28    /// all the children nodes of the call node and at the end it exits.
29    pub fn process_call_node(&mut self, nodes: &[CallTraceNode], idx: usize) {
30        let node = &nodes[idx];
31
32        let func_name = if node.trace.kind.is_any_create() {
33            let contract_name = node
34                .trace
35                .decoded
36                .as_ref()
37                .and_then(|dc| dc.label.as_deref())
38                .unwrap_or("Contract");
39            format!("new {contract_name}")
40        } else {
41            let selector = node
42                .selector()
43                .map(|selector| selector.encode_hex_with_prefix())
44                .unwrap_or_else(|| "fallback".to_string());
45            let signature = node
46                .trace
47                .decoded
48                .as_ref()
49                .and_then(|dc| dc.call_data.as_ref())
50                .map(|dc| &dc.signature)
51                .unwrap_or(&selector);
52
53            if let Some(label) = node.trace.decoded.as_ref().and_then(|dc| dc.label.as_ref()) {
54                format!("{label}.{signature}")
55            } else {
56                signature.clone()
57            }
58        };
59
60        self.fst.enter(func_name, node.trace.gas_used as i64);
61
62        // Track internal function step exits to do in this call context.
63        let mut step_exits = vec![];
64
65        // Process children nodes.
66        for order in &node.ordering {
67            match order {
68                TraceMemberOrder::Call(child_idx) => {
69                    let child_node_idx = node.children[*child_idx];
70                    self.process_call_node(nodes, child_node_idx);
71                }
72                TraceMemberOrder::Step(step_idx) => {
73                    self.exit_previous_steps(&mut step_exits, *step_idx);
74                    self.process_step(&node.trace.steps, *step_idx, &mut step_exits)
75                }
76                TraceMemberOrder::Log(_) => {}
77            }
78        }
79
80        // Exit pending internal function calls if any.
81        for _ in 0..step_exits.len() {
82            self.fst.exit();
83        }
84
85        // Exit from this call context in the folded stack trace.
86        self.fst.exit();
87    }
88
89    /// Creates an entry for an internal function call in the folded stack trace. This method only
90    /// enters the function in the folded stack trace, we cannot exit since we need to exit at a
91    /// future step. Hence, we keep track of the step end index in the `step_exits`.
92    fn process_step(
93        &mut self,
94        steps: &[CallTraceStep],
95        step_idx: usize,
96        step_exits: &mut Vec<usize>,
97    ) {
98        let step = &steps[step_idx];
99        if let Some(decoded_step) = &step.decoded {
100            match decoded_step.as_ref() {
101                DecodedTraceStep::InternalCall(decoded_internal_call, step_end_idx) => {
102                    let gas_used = steps[*step_end_idx].gas_used.saturating_sub(step.gas_used);
103                    self.fst.enter(decoded_internal_call.func_name.clone(), gas_used as i64);
104                    step_exits.push(*step_end_idx);
105                }
106                DecodedTraceStep::Line(_) => {}
107            }
108        }
109    }
110
111    /// Exits all the previous internal calls that should end before starting step_idx.
112    fn exit_previous_steps(&mut self, step_exits: &mut Vec<usize>, step_idx: usize) {
113        let initial_length = step_exits.len();
114        step_exits.retain(|&number| number > step_idx);
115
116        let num_exits = initial_length - step_exits.len();
117        for _ in 0..num_exits {
118            self.fst.exit();
119        }
120    }
121}
122
123/// Helps to translate a function enter-exit flow into a folded stack trace.
124///
125/// Example:
126/// ```solidity
127/// function top() { child_a(); child_b() } // consumes 500 gas
128/// function child_a() {} // consumes 100 gas
129/// function child_b() {} // consumes 200 gas
130/// ```
131///
132/// For execution of the `top` function looks like:
133/// 1. enter `top`
134/// 2. enter `child_a`
135/// 3. exit `child_a`
136/// 4. enter `child_b`
137/// 5. exit `child_b`
138/// 6. exit `top`
139///
140/// The translated folded stack trace lines look like:
141/// 1. top
142/// 2. top;child_a
143/// 3. top;child_b
144///
145/// Including the gas consumed by the function by itself.
146/// 1. top 200 // 500 - 100 - 200
147/// 2. top;child_a 100
148/// 3. top;child_b 200
149#[derive(Debug, Default)]
150pub struct FoldedStackTraceBuilder {
151    /// Trace entries.
152    traces: Vec<TraceEntry>,
153    /// Number of exits to be done before entering a new function.
154    exits: usize,
155}
156
157#[derive(Debug, Default)]
158struct TraceEntry {
159    /// Names of all functions in the call stack of this trace.
160    names: Vec<String>,
161    /// Gas consumed by this function, allowed to be negative due to refunds.
162    gas: i64,
163}
164
165impl FoldedStackTraceBuilder {
166    /// Enter execution of a function call that consumes `gas`.
167    pub fn enter(&mut self, label: String, gas: i64) {
168        let mut names = self.traces.last().map(|entry| entry.names.clone()).unwrap_or_default();
169
170        while self.exits > 0 {
171            names.pop();
172            self.exits -= 1;
173        }
174
175        names.push(label);
176        self.traces.push(TraceEntry { names, gas });
177    }
178
179    /// Exit execution of a function call.
180    pub fn exit(&mut self) {
181        self.exits += 1;
182    }
183
184    /// Returns folded stack trace.
185    pub fn build(mut self) -> Vec<String> {
186        self.subtract_children();
187        self.build_without_subtraction()
188    }
189
190    /// Internal method to build the folded stack trace without subtracting gas consumed by
191    /// the children function calls.
192    fn build_without_subtraction(&mut self) -> Vec<String> {
193        let mut lines = Vec::new();
194        for TraceEntry { names, gas } in &self.traces {
195            lines.push(format!("{} {}", names.join(";"), gas));
196        }
197        lines
198    }
199
200    /// Subtracts gas consumed by the children function calls from the parent function calls.
201    fn subtract_children(&mut self) {
202        // Iterate over each trace to find the children and subtract their values from the parents.
203        for i in 0..self.traces.len() {
204            let (left, right) = self.traces.split_at_mut(i);
205            let TraceEntry { names, gas } = &right[0];
206            if names.len() > 1 {
207                let parent_trace_to_match = &names[..names.len() - 1];
208                for parent in left.iter_mut().rev() {
209                    if parent.names == parent_trace_to_match {
210                        parent.gas -= gas;
211                        break;
212                    }
213                }
214            }
215        }
216    }
217}
218
219mod tests {
220    #[test]
221    fn test_fst_1() {
222        let mut trace = super::FoldedStackTraceBuilder::default();
223        trace.enter("top".to_string(), 500);
224        trace.enter("child_a".to_string(), 100);
225        trace.exit();
226        trace.enter("child_b".to_string(), 200);
227
228        assert_eq!(
229            trace.build_without_subtraction(),
230            vec![
231                "top 500", //
232                "top;child_a 100",
233                "top;child_b 200",
234            ]
235        );
236        assert_eq!(
237            trace.build(),
238            vec![
239                "top 200", // 500 - 100 - 200
240                "top;child_a 100",
241                "top;child_b 200",
242            ]
243        );
244    }
245
246    #[test]
247    fn test_fst_2() {
248        let mut trace = super::FoldedStackTraceBuilder::default();
249        trace.enter("top".to_string(), 500);
250        trace.enter("child_a".to_string(), 300);
251        trace.enter("child_b".to_string(), 100);
252        trace.exit();
253        trace.exit();
254        trace.enter("child_c".to_string(), 100);
255
256        assert_eq!(
257            trace.build_without_subtraction(),
258            vec![
259                "top 500", //
260                "top;child_a 300",
261                "top;child_a;child_b 100",
262                "top;child_c 100",
263            ]
264        );
265
266        assert_eq!(
267            trace.build(),
268            vec![
269                "top 100",         // 500 - 300 - 100
270                "top;child_a 200", // 300 - 100
271                "top;child_a;child_b 100",
272                "top;child_c 100",
273            ]
274        );
275    }
276
277    #[test]
278    fn test_fst_3() {
279        let mut trace = super::FoldedStackTraceBuilder::default();
280        trace.enter("top".to_string(), 1700);
281        trace.enter("child_a".to_string(), 500);
282        trace.exit();
283        trace.enter("child_b".to_string(), 500);
284        trace.enter("child_c".to_string(), 500);
285        trace.exit();
286        trace.exit();
287        trace.exit();
288        trace.enter("top2".to_string(), 1700);
289
290        assert_eq!(
291            trace.build_without_subtraction(),
292            vec![
293                "top 1700", //
294                "top;child_a 500",
295                "top;child_b 500",
296                "top;child_b;child_c 500",
297                "top2 1700",
298            ]
299        );
300
301        assert_eq!(
302            trace.build(),
303            vec![
304                "top 700", //
305                "top;child_a 500",
306                "top;child_b 0",
307                "top;child_b;child_c 500",
308                "top2 1700",
309            ]
310        );
311    }
312}