miden_debug/
flamegraph.rs

1use std::{
2    collections::{BTreeMap, VecDeque},
3    fs::File,
4    io::{BufWriter, Write},
5    path::{Path, PathBuf},
6    sync::Arc,
7};
8
9use miden_assembly::{DefaultSourceManager, SourceManager};
10use miden_assembly_syntax::diagnostics::{IntoDiagnostic, Report, WrapErr};
11use miden_debug_engine::LinkLibrary;
12use miden_processor::ExecutionError;
13
14use crate::{
15    config::{ColorChoice, DebuggerConfig},
16    debug::CallFrame,
17    exec::{DebugExecutor, ExecutionConfig, Executor, ReplaySnapshot},
18    input::InputFile,
19};
20
21/// Folded stack samples keyed by semicolon-separated stack paths.
22pub type Samples = BTreeMap<String, usize>;
23
24/// A collected VM cycle profile that can be rendered as folded stacks or an SVG flamegraph.
25#[derive(Clone, Debug, Default, Eq, PartialEq)]
26pub struct FlamegraphProfile {
27    samples: Samples,
28    total_cycles: usize,
29}
30
31/// The output format selected by [`FlamegraphProfile::write_to_path`].
32#[derive(Clone, Copy, Debug, Eq, PartialEq)]
33pub enum FlamegraphOutput {
34    Svg,
35    FoldedStacks,
36}
37
38impl FlamegraphProfile {
39    /// Execute `executor` to completion, sampling its call stack for every VM cycle.
40    pub fn collect(executor: &mut DebugExecutor) -> Result<Self, ExecutionError> {
41        let mut profile = Self::default();
42
43        loop {
44            if executor.stopped {
45                break;
46            }
47
48            let previous_cycle = executor.cycle;
49            match executor.step() {
50                Ok(_) if executor.cycle > previous_cycle => {
51                    let cycle_delta = executor.cycle - previous_cycle;
52                    profile.record_call_stack(executor.callstack.frames(), cycle_delta);
53                }
54                Ok(_) => {
55                    if executor.stopped {
56                        break;
57                    }
58                }
59                Err(err) => return Err(err),
60            }
61        }
62
63        Ok(profile)
64    }
65
66    /// Replay a recorded execution and collect its cycle-weighted call stacks.
67    pub fn collect_replay(snapshot: ReplaySnapshot) -> Result<Self, ExecutionError> {
68        let ReplaySnapshot {
69            package,
70            stack_inputs,
71            advice_inputs,
72            options,
73            mast_forests,
74            event_log,
75        } = snapshot;
76        let source_manager: Arc<dyn SourceManager> = Arc::new(DefaultSourceManager::default());
77        let executor = Executor::from_config(ExecutionConfig {
78            inputs: stack_inputs,
79            advice_inputs,
80            options,
81        });
82        let mut debug_executor = executor.into_debug_with_replay(
83            package,
84            source_manager,
85            mast_forests,
86            VecDeque::from(event_log),
87        );
88        Self::collect(&mut debug_executor)
89    }
90
91    /// Record `cycles` against a call stack from the debugger engine.
92    pub fn record_call_stack(&mut self, frames: &[CallFrame], cycles: usize) {
93        let path = build_stack_path(frames);
94        self.record_stack_path(path, cycles);
95    }
96
97    /// Record `cycles` against a stack of frame names.
98    ///
99    /// Frame names are sanitized for folded stack output, then joined with `;`.
100    pub fn record_stack<I, S>(&mut self, frames: I, cycles: usize)
101    where
102        I: IntoIterator<Item = S>,
103        S: AsRef<str>,
104    {
105        let path = build_stack_path_from_names(frames);
106        self.record_stack_path(path, cycles);
107    }
108
109    /// Record `cycles` against an already formatted folded-stack path.
110    pub fn record_stack_path(&mut self, stack_path: impl Into<String>, cycles: usize) {
111        if cycles == 0 {
112            return;
113        }
114
115        self.total_cycles += cycles;
116        *self.samples.entry(stack_path.into()).or_default() += cycles;
117    }
118
119    pub fn samples(&self) -> &Samples {
120        &self.samples
121    }
122
123    pub fn total_cycles(&self) -> usize {
124        self.total_cycles
125    }
126
127    pub fn unique_stack_paths(&self) -> usize {
128        self.samples.len()
129    }
130
131    /// Write this profile as an SVG when `path` ends in `.svg`, otherwise as folded stack text.
132    pub fn write_to_path(&self, path: impl AsRef<Path>) -> Result<FlamegraphOutput, Report> {
133        let path = path.as_ref();
134        if is_svg_path(path) {
135            self.write_svg(path)?;
136            Ok(FlamegraphOutput::Svg)
137        } else {
138            self.write_folded_stacks(path)?;
139            Ok(FlamegraphOutput::FoldedStacks)
140        }
141    }
142
143    /// Write this profile in folded stack format.
144    pub fn write_folded_stacks(&self, path: impl AsRef<Path>) -> Result<(), Report> {
145        write_folded_stacks(&self.samples, path.as_ref())
146    }
147
148    /// Render this profile as an SVG flamegraph.
149    pub fn write_svg(&self, path: impl AsRef<Path>) -> Result<(), Report> {
150        generate_svg(&self.samples, path.as_ref())
151    }
152}
153
154#[derive(clap::Args, Debug)]
155#[command(group(
156    clap::ArgGroup::new("execution")
157        .required(true)
158        .args(["input", "replay"])
159))]
160pub struct FlamegraphArgs {
161    /// Specify the path to a Miden program file to execute.
162    #[arg(value_name = "FILE")]
163    pub input: Option<InputFile>,
164    /// Replay a recorded transaction snapshot instead of executing a raw program.
165    #[arg(long, value_name = "FILE")]
166    pub replay: Option<PathBuf>,
167    /// Write the generated flame graph SVG or folded stack text to this path.
168    #[arg(short, long, default_value = "flamegraph.svg")]
169    pub output: PathBuf,
170    /// Specify the path to a file containing program inputs.
171    #[arg(long, value_name = "FILE")]
172    pub inputs: Option<ExecutionConfig>,
173    /// Arguments to encode for the selected program entrypoint.
174    #[arg(last(true), value_name = "ARGV")]
175    pub args: Vec<String>,
176    /// The working directory for execution.
177    #[arg(long, value_name = "DIR", help_heading = "Execution")]
178    pub working_dir: Option<PathBuf>,
179    /// The path to the root directory of the current Miden toolchain.
180    #[arg(
181        long,
182        value_name = "DIR",
183        env = "MIDEN_SYSROOT",
184        help_heading = "Linker"
185    )]
186    pub sysroot: Option<PathBuf>,
187    /// Specify the function to call as the entrypoint for the program.
188    #[arg(long, help_heading = "Execution")]
189    pub entrypoint: Option<String>,
190    /// Specify one or more search paths for link libraries requested via `-l`.
191    #[arg(
192        long = "search-path",
193        short = 'L',
194        value_name = "PATH",
195        help_heading = "Linker"
196    )]
197    pub search_path: Vec<PathBuf>,
198    /// Link compiled projects to the specified library NAME.
199    #[arg(
200        long = "link-library",
201        short = 'l',
202        value_name = "[KIND=]NAME",
203        value_delimiter = ',',
204        next_line_help(true),
205        help_heading = "Linker"
206    )]
207    pub link_libraries: Vec<LinkLibrary>,
208}
209
210impl FlamegraphArgs {
211    fn into_debugger_config(self) -> DebuggerConfig {
212        DebuggerConfig {
213            input: self.input,
214            inputs: self.inputs,
215            args: self.args,
216            working_dir: self.working_dir,
217            sysroot: self.sysroot,
218            color: ColorChoice::Auto,
219            entrypoint: self.entrypoint,
220            #[cfg(all(feature = "dap", feature = "tui"))]
221            dap_connect: None,
222            #[cfg(feature = "dap")]
223            start_debug_adapter: None,
224            source_path_prefixes: Vec::new(),
225            search_path: self.search_path,
226            link_libraries: self.link_libraries,
227            repl: false,
228            commands: None,
229            #[cfg(feature = "tui")]
230            replay: None,
231            #[cfg(feature = "python")]
232            no_user_python_init: true,
233            profiler_cli_args: Default::default(), // disable profiling
234        }
235    }
236}
237
238pub fn run(args: FlamegraphArgs) -> Result<(), Report> {
239    let output = args.output.clone();
240    let replay = args.replay.clone();
241    if let Some(replay) = replay {
242        let snapshot = ReplaySnapshot::read_from_file(&replay)
243            .map_err(|err| Report::msg(format!("failed to read {}: {err}", replay.display())))?;
244        let profile = FlamegraphProfile::collect_replay(snapshot)
245            .map_err(|err| Report::msg(format!("replay execution failed: {err}")))?;
246        report_and_write_profile(&profile, &output)?;
247        return Ok(());
248    }
249
250    let mut config = args.into_debugger_config();
251    ensure_working_dir(&mut config)?;
252
253    let source_manager: Arc<dyn SourceManager> = Arc::new(DefaultSourceManager::default());
254    let mut executor =
255        crate::program_loader::load_debug_executor(&config, source_manager, "flamegraph")?.executor;
256
257    let profile = match FlamegraphProfile::collect(&mut executor) {
258        Ok(profile) => profile,
259        Err(err) => {
260            return Err(Report::msg(format!(
261                "program execution failed at cycle {}: {err}",
262                executor.cycle
263            )));
264        }
265    };
266
267    report_and_write_profile(&profile, &output)?;
268
269    Ok(())
270}
271
272fn report_and_write_profile(profile: &FlamegraphProfile, output: &Path) -> Result<(), Report> {
273    eprintln!(
274        "Executed {} cycles across {} unique stack paths",
275        profile.total_cycles(),
276        profile.unique_stack_paths()
277    );
278    profile.write_to_path(output)?;
279    Ok(())
280}
281
282fn ensure_working_dir(config: &mut DebuggerConfig) -> Result<(), Report> {
283    if config.working_dir.is_none() {
284        let cwd = std::env::current_dir()
285            .into_diagnostic()
286            .wrap_err("could not read current working directory")?;
287        config.working_dir = Some(cwd);
288    }
289
290    Ok(())
291}
292
293fn build_stack_path(frames: &[CallFrame]) -> String {
294    build_stack_path_from_names(frames.iter().filter_map(|frame| frame.procedure("")))
295}
296
297fn build_stack_path_from_names<I, S>(frames: I) -> String
298where
299    I: IntoIterator<Item = S>,
300    S: AsRef<str>,
301{
302    let mut path = String::new();
303    for frame in frames {
304        append_frame(&mut path, frame.as_ref());
305    }
306
307    if path.is_empty() {
308        "[unknown]".to_string()
309    } else {
310        path
311    }
312}
313
314fn append_frame(path: &mut String, name: &str) {
315    if !path.is_empty() {
316        path.push(';');
317    }
318    append_sanitized_frame(path, name);
319}
320
321fn append_sanitized_frame(path: &mut String, name: &str) {
322    for ch in name.chars() {
323        match ch {
324            ';' => path.push(':'),
325            '\n' | '\r' => path.push(' '),
326            _ => path.push(ch),
327        }
328    }
329}
330
331fn is_svg_path(path: &Path) -> bool {
332    path.extension()
333        .and_then(|ext| ext.to_str())
334        .is_some_and(|ext| ext.eq_ignore_ascii_case("svg"))
335}
336
337fn write_folded_stacks(samples: &Samples, path: &Path) -> Result<(), Report> {
338    let file = File::create(path).into_diagnostic()?;
339    let mut writer = BufWriter::new(file);
340    for (stack, count) in samples {
341        writeln!(writer, "{stack} {count}").into_diagnostic()?;
342    }
343    writer.flush().into_diagnostic()?;
344
345    eprintln!("Wrote folded stacks to {}", path.display());
346    Ok(())
347}
348
349fn generate_svg(samples: &Samples, path: &Path) -> Result<(), Report> {
350    let input = samples
351        .iter()
352        .map(|(stack, count)| format!("{stack} {count}"))
353        .collect::<Vec<_>>()
354        .join("\n");
355
356    let mut opts = inferno::flamegraph::Options::default();
357    opts.title = "Miden VM Flame Graph (cycles)".to_string();
358    opts.count_name = "cycles".to_string();
359
360    let file = File::create(path).into_diagnostic()?;
361    let mut writer = BufWriter::new(file);
362    inferno::flamegraph::from_reader(&mut opts, input.as_bytes(), &mut writer).into_diagnostic()?;
363    writer.flush().into_diagnostic()?;
364
365    eprintln!("Wrote flame graph to {}", path.display());
366    Ok(())
367}
368
369#[cfg(test)]
370mod tests {
371    use super::FlamegraphProfile;
372
373    #[test]
374    fn record_stack_sanitizes_folded_stack_frames() {
375        let mut profile = FlamegraphProfile::default();
376
377        profile.record_stack(["root;proc", "child\nproc"], 3);
378        profile.record_stack(["root;proc", "child\nproc"], 2);
379        profile.record_stack(["ignored"], 0);
380
381        assert_eq!(profile.total_cycles(), 5);
382        assert_eq!(profile.unique_stack_paths(), 1);
383        assert_eq!(profile.samples().get("root:proc;child proc"), Some(&5));
384        assert!(!profile.samples().contains_key("ignored"));
385    }
386}
miden_debug/flamegraph.rs

miden_debug/
flamegraph.rs
