cu_profiler_core/parser/
cpi_tree.rs

1//! Reconstructs the CPI call tree from a lexed event stream.
2//!
3//! The builder is tolerant: mismatched or missing success/failure lines never
4//! panic; nodes simply close on a best-effort basis and the resulting tree is
5//! still well-formed.
6
7use serde::{Deserialize, Serialize};
8
9use crate::parser::solana_logs::LogEvent;
10use crate::program_registry::ProgramRegistry;
11
12/// Outcome of a single program invocation.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "lowercase")]
15pub enum NodeStatus {
16    /// Reported `success`.
17    Success,
18    /// Reported `failed`.
19    Failed,
20    /// No terminal line was seen.
21    Unknown,
22}
23
24/// A node in the CPI call tree.
25#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
26pub struct CallNode {
27    /// Program ID (or `"transaction"` for the synthetic root).
28    pub program_id: String,
29    /// Resolved display label, if known.
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub label: Option<String>,
32    /// Invoke depth (0 for the root, 1 for the entrypoint program).
33    pub depth: u32,
34    /// CU consumed by this invocation, if reported.
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub units_consumed: Option<u64>,
37    /// Terminal status.
38    pub status: NodeStatus,
39    /// Log messages emitted directly under this invocation.
40    #[serde(default, skip_serializing_if = "Vec::is_empty")]
41    pub logs: Vec<String>,
42    /// Nested invocations.
43    #[serde(default, skip_serializing_if = "Vec::is_empty")]
44    pub children: Vec<CallNode>,
45}
46
47impl CallNode {
48    fn new(program_id: String, depth: u32, registry: &ProgramRegistry) -> Self {
49        let label = registry.label(&program_id).map(str::to_string);
50        Self {
51            program_id,
52            label,
53            depth,
54            units_consumed: None,
55            status: NodeStatus::Unknown,
56            logs: Vec::new(),
57            children: Vec::new(),
58        }
59    }
60
61    /// The synthetic transaction root.
62    fn root() -> Self {
63        Self {
64            program_id: "transaction".to_string(),
65            label: Some("root transaction".to_string()),
66            depth: 0,
67            units_consumed: None,
68            status: NodeStatus::Unknown,
69            logs: Vec::new(),
70            children: Vec::new(),
71        }
72    }
73
74    /// Total number of program invocations in the subtree (excluding the root).
75    #[must_use]
76    pub fn invocation_count(&self) -> u32 {
77        self.children.iter().map(|c| 1 + c.invocation_count()).sum()
78    }
79
80    /// Maximum invoke depth in the subtree.
81    #[must_use]
82    pub fn max_depth(&self) -> u32 {
83        self.children
84            .iter()
85            .map(CallNode::max_depth)
86            .max()
87            .map_or(self.depth, |d| d.max(self.depth))
88    }
89
90    /// Number of cross-program invocations (depth >= 2).
91    #[must_use]
92    pub fn cpi_count(&self) -> u32 {
93        let here = u32::from(self.depth >= 2);
94        here + self.children.iter().map(CallNode::cpi_count).sum::<u32>()
95    }
96}
97
98/// Maximum CPI nesting the tree will represent. Real Solana caps CPI depth at a
99/// handful of levels; this generous bound prevents adversarial logs (tens of
100/// thousands of unclosed `invoke` lines) from building a tree so deep that
101/// serializing or traversing it overflows the stack. Invocations beyond it are
102/// flattened (attached at the cap, not nested deeper).
103pub const MAX_DEPTH: usize = 64;
104
105/// Build a call tree from a lexed event stream.
106#[must_use]
107pub fn build(events: &[LogEvent], registry: &ProgramRegistry) -> CallNode {
108    let mut stack: Vec<CallNode> = vec![CallNode::root()];
109
110    for event in events {
111        match event {
112            LogEvent::Invoke { program_id, depth } => {
113                let node = CallNode::new(program_id.clone(), *depth, registry);
114                // Cap nesting: beyond MAX_DEPTH, attach as a leaf of the deepest
115                // open frame instead of opening a new one. Closes still balance
116                // against the real open frames; the tree depth stays bounded.
117                if stack.len() <= MAX_DEPTH {
118                    stack.push(node);
119                } else if let Some(top) = stack.last_mut() {
120                    top.children.push(node);
121                }
122            }
123            LogEvent::Consumed {
124                program_id, used, ..
125            } => {
126                // Assign only on an exact program-ID match. In real Solana logs
127                // a program's `consumed` line is emitted while that program is
128                // the open frame, so this always matches; refusing to fall back
129                // to "any open frame" prevents silent CU misattribution when
130                // logs are malformed or out of order.
131                if let Some(top) = stack.last_mut() {
132                    if &top.program_id == program_id {
133                        top.units_consumed = Some(*used);
134                    }
135                }
136            }
137            LogEvent::Success { .. } => close_top(&mut stack, NodeStatus::Success),
138            LogEvent::Failed { .. } => close_top(&mut stack, NodeStatus::Failed),
139            LogEvent::Log { message } => push_log(&mut stack, message.clone()),
140            LogEvent::ScopeBegin { name, .. } => {
141                push_log(&mut stack, format!("scope-begin: {name}"));
142            }
143            LogEvent::ScopeEnd { name, .. } => push_log(&mut stack, format!("scope-end: {name}")),
144            LogEvent::ScopePoint { name, .. } => {
145                push_log(&mut stack, format!("scope-point: {name}"));
146            }
147            LogEvent::Raw(line) => push_log(&mut stack, line.clone()),
148        }
149    }
150
151    // Anything left open closes into the root, preserving partial trees.
152    while stack.len() > 1 {
153        close_top(&mut stack, NodeStatus::Unknown);
154    }
155    stack.pop().unwrap_or_else(CallNode::root)
156}
157
158fn close_top(stack: &mut Vec<CallNode>, status: NodeStatus) {
159    if stack.len() <= 1 {
160        return; // never pop the root
161    }
162    let mut node = stack.pop().expect("len > 1 checked");
163    if node.status == NodeStatus::Unknown {
164        node.status = status;
165    }
166    if let Some(parent) = stack.last_mut() {
167        parent.children.push(node);
168    }
169}
170
171fn push_log(stack: &mut [CallNode], message: String) {
172    if let Some(top) = stack.last_mut() {
173        top.logs.push(message);
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::parser::solana_logs::lex;
181
182    fn tree_from(lines: &[&str]) -> CallNode {
183        let owned: Vec<String> = lines.iter().map(|s| (*s).to_string()).collect();
184        let lexed = lex(&owned);
185        let events: Vec<LogEvent> = lexed.events().cloned().collect();
186        build(&events, &ProgramRegistry::with_builtins())
187    }
188
189    #[test]
190    fn builds_nested_cpi_tree() {
191        let tree = tree_from(&[
192            "Program User111 invoke [1]",
193            "Program TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA invoke [2]",
194            "Program TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA consumed 3000 of 197000 compute units",
195            "Program TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA success",
196            "Program User111 consumed 12000 of 200000 compute units",
197            "Program User111 success",
198        ]);
199        assert_eq!(tree.children.len(), 1);
200        let user = &tree.children[0];
201        assert_eq!(user.program_id, "User111");
202        assert_eq!(user.units_consumed, Some(12000));
203        assert_eq!(user.children.len(), 1);
204        assert_eq!(user.children[0].label.as_deref(), Some("SPL Token"));
205        assert_eq!(tree.cpi_count(), 1);
206        assert_eq!(tree.max_depth(), 2);
207        assert_eq!(tree.invocation_count(), 2);
208    }
209
210    #[test]
211    fn deep_invoke_chain_is_capped_not_overflowing() {
212        // Tens of thousands of unclosed invokes would build a tree deep enough
213        // to overflow the stack on serialize/traverse; the cap flattens it.
214        let lines: Vec<String> = (0..50_000)
215            .map(|i| format!("Program P{i} invoke [{}]", i + 1))
216            .collect();
217        let lexed = lex(&lines);
218        let events: Vec<LogEvent> = lexed.events().cloned().collect();
219        let tree = build(&events, &ProgramRegistry::with_builtins());
220        // These recursive traversals must not overflow (bounded by MAX_DEPTH).
221        assert!(tree.invocation_count() >= 50_000);
222        assert!(tree.max_depth() >= 1);
223        // And the structure is genuinely shallow now.
224        fn nesting(n: &CallNode) -> usize {
225            1 + n.children.iter().map(nesting).max().unwrap_or(0)
226        }
227        // root + MAX_DEPTH open frames + one flattened-leaf level.
228        assert!(nesting(&tree) <= MAX_DEPTH + 2);
229    }
230
231    #[test]
232    fn unterminated_invoke_does_not_panic() {
233        let tree = tree_from(&["Program User111 invoke [1]"]);
234        assert_eq!(tree.children.len(), 1);
235        assert_eq!(tree.children[0].status, NodeStatus::Unknown);
236    }
237
238    #[test]
239    fn consumed_attributes_only_to_matching_program() {
240        // A `consumed` line for a program that is NOT the open frame must not be
241        // misattributed to whatever frame happens to be on top of the stack.
242        let tree = tree_from(&[
243            "Program A111 invoke [1]",
244            "Program B222 invoke [2]",
245            "Program A111 consumed 5000 of 200000 compute units", // B is open, not A
246            "Program B222 consumed 3000 of 197000 compute units",
247            "Program B222 success",
248            "Program A111 success",
249        ]);
250        let a = &tree.children[0];
251        let b = &a.children[0];
252        assert_eq!(b.program_id, "B222");
253        // B keeps its own figure; the stray "A consumed" did not land on B.
254        assert_eq!(b.units_consumed, Some(3000));
255    }
256}
cu_profiler_core/parser/cpi_tree.rs

cu_profiler_core/parser/
cpi_tree.rs
