Skip to main content

stylus_trace_core/aggregator/
stack_builder.rs

1//! Build collapsed stack format from parsed trace data.
2//!
3//! Collapsed stacks are the input format for flamegraph generation.
4//! Format: "parent;child;grandchild weight"
5//!
6//! Example: "main;execute_tx;storage_read 1000"
7//! This means: main called execute_tx which called storage_read, consuming 1000 gas.
8
9use crate::parser::{HostIoType, ParsedTrace};
10use log::debug;
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13
14/// A single collapsed stack entry
15///
16/// **Public** - used by flamegraph generator
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct CollapsedStack {
19    /// Stack trace as semicolon-separated string
20    pub stack: String,
21
22    /// Weight (gas consumed by this stack)
23    pub weight: u64,
24
25    /// Last Program Counter / Offset associated with this stack
26    pub last_pc: Option<u64>,
27}
28
29impl CollapsedStack {
30    /// Create a new collapsed stack
31    ///
32    /// **Public** - constructor
33    pub fn new(stack: String, weight: u64, last_pc: Option<u64>) -> Self {
34        Self {
35            stack,
36            weight,
37            last_pc,
38        }
39    }
40}
41
42/// Build collapsed stacks from parsed trace
43///
44/// **Public** - main entry point for stack building
45///
46/// # Arguments
47/// * `parsed_trace` - Parsed trace data from parser
48///
49/// # Returns
50/// Vector of collapsed stacks, one per unique execution path
51///
52/// # Algorithm
53/// 1. Walk through execution steps
54/// 2. Track call stack depth
55/// 3. Build stack strings for each gas-consuming operation
56/// 4. Aggregate by unique stack (sum weights)
57pub fn build_collapsed_stacks(parsed_trace: &ParsedTrace) -> Vec<CollapsedStack> {
58    debug!(
59        "Building collapsed stacks from {} execution steps",
60        parsed_trace.execution_steps.len()
61    );
62
63    // Map to aggregate stacks: stack_string -> (total_weight, last_pc)
64    let mut stack_map: HashMap<String, (u64, u64)> = HashMap::new();
65
66    // Current call stack (tracks function hierarchy)
67    let mut call_stack: Vec<String> = Vec::new();
68
69    // Process each execution step
70    for step in &parsed_trace.execution_steps {
71        // Get operation name and map to HostIO name if it's an opcode
72        let raw_op = step
73            .function
74            .as_deref()
75            .or(step.op.as_deref())
76            .unwrap_or("unknown");
77
78        // Handle formats like "call;SSTORE"
79        let op_part = raw_op.split(';').next_back().unwrap_or(raw_op);
80
81        let operation = HostIoType::from_opcode(op_part)
82            .map(map_hostio_to_label)
83            .unwrap_or(raw_op);
84
85        // Handle depth changes properly
86        let current_depth = step.depth as usize;
87
88        // If depth decreased, we returned from function calls
89        if current_depth < call_stack.len() {
90            call_stack.truncate(current_depth);
91        }
92
93        // If depth increased, we entered a new call
94        // (This happens if we missed some steps or have shallow tracing)
95        while call_stack.len() < current_depth {
96            call_stack.push("call".to_string());
97        }
98
99        // Build the full stack string with current operation
100        let stack_str = if call_stack.is_empty() {
101            operation.to_string()
102        } else {
103            format!("{};{}", call_stack.join(";"), operation)
104        };
105
106        // Accumulate all gas costs
107        let entry = stack_map.entry(stack_str).or_insert((0, 0));
108        entry.0 += step.gas_cost;
109        entry.1 = step.pc;
110    }
111
112    // Convert map to vector and sort by weight (descending)
113    let mut stacks: Vec<CollapsedStack> = stack_map
114        .into_iter()
115        .map(|(stack, (weight, pc))| CollapsedStack::new(stack, weight, Some(pc)))
116        .collect();
117
118    stacks.sort_by(|a, b| b.weight.cmp(&a.weight));
119    debug!("Built {} unique collapsed stacks", stacks.len());
120
121    stacks
122}
123
124/// Map HostIO type to human-readable label
125pub fn map_hostio_to_label(io_type: HostIoType) -> &'static str {
126    match io_type {
127        HostIoType::StorageLoad => "storage_load_bytes32",
128        HostIoType::StorageStore => "storage_store_bytes32",
129        HostIoType::StorageFlush => "storage_flush_cache",
130        HostIoType::StorageCache => "storage_cache",
131        HostIoType::Call => "call",
132        HostIoType::StaticCall => "staticcall",
133        HostIoType::DelegateCall => "delegatecall",
134        HostIoType::Create => "create",
135        HostIoType::Log => "emit_log",
136        HostIoType::SelfDestruct => "selfdestruct",
137        HostIoType::AccountBalance => "account_balance",
138        HostIoType::BlockHash => "block_hash",
139        HostIoType::NativeKeccak256 => "native_keccak256",
140        HostIoType::ReadArgs => "read_args",
141        HostIoType::WriteResult => "write_result",
142        HostIoType::MsgValue => "msg_value",
143        HostIoType::MsgSender => "msg_sender",
144        HostIoType::MsgReentrant => "msg_reentrant",
145        HostIoType::Other => "other",
146    }
147}