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/// Build a call tree from a lexed event stream.
99#[must_use]
100pub fn build(events: &[LogEvent], registry: &ProgramRegistry) -> CallNode {
101    let mut stack: Vec<CallNode> = vec![CallNode::root()];
102
103    for event in events {
104        match event {
105            LogEvent::Invoke { program_id, depth } => {
106                stack.push(CallNode::new(program_id.clone(), *depth, registry));
107            }
108            LogEvent::Consumed {
109                program_id, used, ..
110            } => {
111                // Assign only on an exact program-ID match. In real Solana logs
112                // a program's `consumed` line is emitted while that program is
113                // the open frame, so this always matches; refusing to fall back
114                // to "any open frame" prevents silent CU misattribution when
115                // logs are malformed or out of order.
116                if let Some(top) = stack.last_mut() {
117                    if &top.program_id == program_id {
118                        top.units_consumed = Some(*used);
119                    }
120                }
121            }
122            LogEvent::Success { .. } => close_top(&mut stack, NodeStatus::Success),
123            LogEvent::Failed { .. } => close_top(&mut stack, NodeStatus::Failed),
124            LogEvent::Log { message } => push_log(&mut stack, message.clone()),
125            LogEvent::ScopeBegin { name, .. } => {
126                push_log(&mut stack, format!("scope-begin: {name}"));
127            }
128            LogEvent::ScopeEnd { name, .. } => push_log(&mut stack, format!("scope-end: {name}")),
129            LogEvent::ScopePoint { name, .. } => {
130                push_log(&mut stack, format!("scope-point: {name}"));
131            }
132            LogEvent::Raw(line) => push_log(&mut stack, line.clone()),
133        }
134    }
135
136    // Anything left open closes into the root, preserving partial trees.
137    while stack.len() > 1 {
138        close_top(&mut stack, NodeStatus::Unknown);
139    }
140    stack.pop().unwrap_or_else(CallNode::root)
141}
142
143fn close_top(stack: &mut Vec<CallNode>, status: NodeStatus) {
144    if stack.len() <= 1 {
145        return; // never pop the root
146    }
147    let mut node = stack.pop().expect("len > 1 checked");
148    if node.status == NodeStatus::Unknown {
149        node.status = status;
150    }
151    if let Some(parent) = stack.last_mut() {
152        parent.children.push(node);
153    }
154}
155
156fn push_log(stack: &mut [CallNode], message: String) {
157    if let Some(top) = stack.last_mut() {
158        top.logs.push(message);
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use crate::parser::solana_logs::lex;
166
167    fn tree_from(lines: &[&str]) -> CallNode {
168        let owned: Vec<String> = lines.iter().map(|s| (*s).to_string()).collect();
169        let lexed = lex(&owned);
170        let events: Vec<LogEvent> = lexed.events().cloned().collect();
171        build(&events, &ProgramRegistry::with_builtins())
172    }
173
174    #[test]
175    fn builds_nested_cpi_tree() {
176        let tree = tree_from(&[
177            "Program User111 invoke [1]",
178            "Program TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA invoke [2]",
179            "Program TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA consumed 3000 of 197000 compute units",
180            "Program TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA success",
181            "Program User111 consumed 12000 of 200000 compute units",
182            "Program User111 success",
183        ]);
184        assert_eq!(tree.children.len(), 1);
185        let user = &tree.children[0];
186        assert_eq!(user.program_id, "User111");
187        assert_eq!(user.units_consumed, Some(12000));
188        assert_eq!(user.children.len(), 1);
189        assert_eq!(user.children[0].label.as_deref(), Some("SPL Token"));
190        assert_eq!(tree.cpi_count(), 1);
191        assert_eq!(tree.max_depth(), 2);
192        assert_eq!(tree.invocation_count(), 2);
193    }
194
195    #[test]
196    fn unterminated_invoke_does_not_panic() {
197        let tree = tree_from(&["Program User111 invoke [1]"]);
198        assert_eq!(tree.children.len(), 1);
199        assert_eq!(tree.children[0].status, NodeStatus::Unknown);
200    }
201
202    #[test]
203    fn consumed_attributes_only_to_matching_program() {
204        // A `consumed` line for a program that is NOT the open frame must not be
205        // misattributed to whatever frame happens to be on top of the stack.
206        let tree = tree_from(&[
207            "Program A111 invoke [1]",
208            "Program B222 invoke [2]",
209            "Program A111 consumed 5000 of 200000 compute units", // B is open, not A
210            "Program B222 consumed 3000 of 197000 compute units",
211            "Program B222 success",
212            "Program A111 success",
213        ]);
214        let a = &tree.children[0];
215        let b = &a.children[0];
216        assert_eq!(b.program_id, "B222");
217        // B keeps its own figure; the stray "A consumed" did not land on B.
218        assert_eq!(b.units_consumed, Some(3000));
219    }
220}
cu_profiler_core/parser/cpi_tree.rs

cu_profiler_core/parser/
cpi_tree.rs
