miden_debug/
flamegraph.rs

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

miden_debug/
flamegraph.rs
