miden_debug/
flamegraph.rs

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

miden_debug/
flamegraph.rs
