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::{
11    Library,
12    diagnostics::{IntoDiagnostic, Report, WrapErr},
13};
14use miden_core::serde::Deserializable;
15use miden_processor::{ExecutionError, StackInputs};
16
17use crate::{
18    config::{ColorChoice, DebuggerConfig},
19    debug::CallFrame,
20    exec::{DebugExecutor, ExecutionConfig, Executor},
21    felt::Felt,
22    input::InputFile,
23    linker::LinkLibrary,
24};
25
26/// Folded stack samples keyed by semicolon-separated stack paths.
27pub type Samples = BTreeMap<String, usize>;
28
29/// A collected VM cycle profile that can be rendered as folded stacks or an SVG flamegraph.
30#[derive(Clone, Debug, Default, Eq, PartialEq)]
31pub struct FlamegraphProfile {
32    samples: Samples,
33    total_cycles: usize,
34}
35
36/// The output format selected by [`FlamegraphProfile::write_to_path`].
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
38pub enum FlamegraphOutput {
39    Svg,
40    FoldedStacks,
41}
42
43impl FlamegraphProfile {
44    /// Execute `executor` to completion, sampling its call stack for every VM cycle.
45    pub fn collect(executor: &mut DebugExecutor) -> Result<Self, ExecutionError> {
46        let mut profile = Self::default();
47
48        loop {
49            if executor.stopped {
50                break;
51            }
52
53            let previous_cycle = executor.cycle;
54            match executor.step() {
55                Ok(_) if executor.cycle > previous_cycle => {
56                    let cycle_delta = executor.cycle - previous_cycle;
57                    profile.record_call_stack(executor.callstack.frames(), cycle_delta);
58                }
59                Ok(_) => {
60                    if executor.stopped {
61                        break;
62                    }
63                }
64                Err(err) => return Err(err),
65            }
66        }
67
68        Ok(profile)
69    }
70
71    /// Record `cycles` against a call stack from the debugger engine.
72    pub fn record_call_stack(&mut self, frames: &[CallFrame], cycles: usize) {
73        let path = build_stack_path(frames);
74        self.record_stack_path(path, cycles);
75    }
76
77    /// Record `cycles` against a stack of frame names.
78    ///
79    /// Frame names are sanitized for folded stack output, then joined with `;`.
80    pub fn record_stack<I, S>(&mut self, frames: I, cycles: usize)
81    where
82        I: IntoIterator<Item = S>,
83        S: AsRef<str>,
84    {
85        let path = build_stack_path_from_names(frames);
86        self.record_stack_path(path, cycles);
87    }
88
89    /// Record `cycles` against an already formatted folded-stack path.
90    pub fn record_stack_path(&mut self, stack_path: impl Into<String>, cycles: usize) {
91        if cycles == 0 {
92            return;
93        }
94
95        self.total_cycles += cycles;
96        *self.samples.entry(stack_path.into()).or_default() += cycles;
97    }
98
99    pub fn samples(&self) -> &Samples {
100        &self.samples
101    }
102
103    pub fn total_cycles(&self) -> usize {
104        self.total_cycles
105    }
106
107    pub fn unique_stack_paths(&self) -> usize {
108        self.samples.len()
109    }
110
111    /// Write this profile as an SVG when `path` ends in `.svg`, otherwise as folded stack text.
112    pub fn write_to_path(&self, path: impl AsRef<Path>) -> Result<FlamegraphOutput, Report> {
113        let path = path.as_ref();
114        if is_svg_path(path) {
115            self.write_svg(path)?;
116            Ok(FlamegraphOutput::Svg)
117        } else {
118            self.write_folded_stacks(path)?;
119            Ok(FlamegraphOutput::FoldedStacks)
120        }
121    }
122
123    /// Write this profile in folded stack format.
124    pub fn write_folded_stacks(&self, path: impl AsRef<Path>) -> Result<(), Report> {
125        write_folded_stacks(&self.samples, path.as_ref())
126    }
127
128    /// Render this profile as an SVG flamegraph.
129    pub fn write_svg(&self, path: impl AsRef<Path>) -> Result<(), Report> {
130        generate_svg(&self.samples, path.as_ref())
131    }
132}
133
134#[derive(clap::Args, Debug)]
135pub struct FlamegraphArgs {
136    /// Specify the path to a Miden program file to execute.
137    #[arg(value_name = "FILE")]
138    pub input: InputFile,
139    /// Write the generated flame graph SVG or folded stack text to this path.
140    #[arg(short, long, default_value = "flamegraph.svg")]
141    pub output: PathBuf,
142    /// Specify the path to a file containing program inputs.
143    #[arg(long, value_name = "FILE")]
144    pub inputs: Option<ExecutionConfig>,
145    /// Arguments to place on the operand stack before calling the program entrypoint.
146    #[arg(last(true), value_name = "ARGV")]
147    pub args: Vec<Felt>,
148    /// The working directory for execution.
149    #[arg(long, value_name = "DIR", help_heading = "Execution")]
150    pub working_dir: Option<PathBuf>,
151    /// The path to the root directory of the current Miden toolchain.
152    #[arg(
153        long,
154        value_name = "DIR",
155        env = "MIDEN_SYSROOT",
156        help_heading = "Linker"
157    )]
158    pub sysroot: Option<PathBuf>,
159    /// Specify the function to call as the entrypoint for the program.
160    #[arg(long, help_heading = "Execution")]
161    pub entrypoint: Option<String>,
162    /// Specify one or more search paths for link libraries requested via `-l`.
163    #[arg(
164        long = "search-path",
165        short = 'L',
166        value_name = "PATH",
167        help_heading = "Linker"
168    )]
169    pub search_path: Vec<PathBuf>,
170    /// Link compiled projects to the specified library NAME.
171    #[arg(
172        long = "link-library",
173        short = 'l',
174        value_name = "[KIND=]NAME",
175        value_delimiter = ',',
176        next_line_help(true),
177        help_heading = "Linker"
178    )]
179    pub link_libraries: Vec<LinkLibrary>,
180    /// Source path prefixes used by the compiler's `-Zremap-path-prefix` option.
181    ///
182    /// When debug info stores trimmed source paths, DAP clients may still send
183    /// absolute editor paths. These prefixes provide an explicit mapping between
184    /// the two forms.
185    #[cfg(feature = "dap")]
186    #[cfg_attr(
187        feature = "tui",
188        arg(
189            long = "source-path-prefix",
190            alias = "trim-path-prefix",
191            value_name = "PATH",
192            help_heading = "Debugging"
193        )
194    )]
195    pub source_path_prefixes: Vec<PathBuf>,
196}
197
198impl FlamegraphArgs {
199    fn into_debugger_config(self) -> DebuggerConfig {
200        DebuggerConfig {
201            input: Some(self.input),
202            inputs: self.inputs,
203            args: self.args,
204            working_dir: self.working_dir,
205            sysroot: self.sysroot,
206            color: ColorChoice::Auto,
207            entrypoint: self.entrypoint,
208            #[cfg(feature = "dap")]
209            dap_connect: None,
210            #[cfg(feature = "dap")]
211            start_debug_adapter: None,
212            #[cfg(feature = "dap")]
213            source_path_prefixes: self.source_path_prefixes,
214            search_path: self.search_path,
215            link_libraries: self.link_libraries,
216            repl: false,
217        }
218    }
219}
220
221pub fn run(args: FlamegraphArgs) -> Result<(), Report> {
222    let output = args.output.clone();
223    let mut config = args.into_debugger_config();
224    ensure_working_dir(&mut config)?;
225
226    let source_manager: Arc<dyn SourceManager> = Arc::new(DefaultSourceManager::default());
227    let mut inputs = execution_inputs(&config)?;
228    let args = inputs.inputs.iter().copied().collect::<Vec<_>>();
229    let package = load_package(&config)?;
230
231    let libs = load_libraries(&config, source_manager.clone())?;
232
233    let mut executor = Executor::new(args);
234    for lib in libs.iter() {
235        executor.register_library_dependency(lib.clone());
236        executor.with_library(lib.clone());
237    }
238    executor.with_dependencies(package.manifest.dependencies())?;
239    executor.with_advice_inputs(core::mem::take(&mut inputs.advice_inputs));
240
241    let program = package.unwrap_program();
242    let mut executor = executor.into_debug(&program, source_manager);
243
244    let profile = match FlamegraphProfile::collect(&mut executor) {
245        Ok(profile) => profile,
246        Err(err) => {
247            return Err(Report::msg(format!(
248                "program execution failed at cycle {}: {err}",
249                executor.cycle
250            )));
251        }
252    };
253
254    eprintln!(
255        "Executed {} cycles across {} unique stack paths",
256        profile.total_cycles(),
257        profile.unique_stack_paths()
258    );
259
260    profile.write_to_path(&output)?;
261
262    Ok(())
263}
264
265fn ensure_working_dir(config: &mut DebuggerConfig) -> Result<(), Report> {
266    if config.working_dir.is_none() {
267        let cwd = std::env::current_dir()
268            .into_diagnostic()
269            .wrap_err("could not read current working directory")?;
270        config.working_dir = Some(cwd);
271    }
272
273    Ok(())
274}
275
276fn execution_inputs(config: &DebuggerConfig) -> Result<ExecutionConfig, Report> {
277    let mut inputs = config.inputs.clone().unwrap_or_default();
278    if !config.args.is_empty() {
279        // CLI args model sequential pushes, but StackInputs expects the top element first.
280        let args = config.args.iter().rev().map(|felt| felt.0).collect::<Vec<_>>();
281        inputs.inputs = StackInputs::new(&args).into_diagnostic()?;
282    }
283
284    Ok(inputs)
285}
286
287fn load_libraries(
288    config: &DebuggerConfig,
289    source_manager: Arc<dyn SourceManager>,
290) -> Result<Vec<Arc<Library>>, Report> {
291    let mut libs = Vec::with_capacity(config.link_libraries.len());
292    for link_library in config.link_libraries.iter() {
293        log::debug!(target: "flamegraph", "loading link library {}", link_library.name());
294        libs.push(link_library.load(config, source_manager.clone())?);
295    }
296
297    if let Some(toolchain_dir) = config.toolchain_dir() {
298        libs.extend(load_sysroot_libs(&toolchain_dir)?);
299    }
300
301    Ok(libs)
302}
303
304fn load_sysroot_libs(toolchain_dir: &Path) -> Result<Vec<Arc<Library>>, Report> {
305    let mut libs = Vec::new();
306
307    let entries = match std::fs::read_dir(toolchain_dir) {
308        Ok(entries) => entries,
309        Err(_) => {
310            log::debug!(target: "flamegraph", "could not read sysroot directory: {}", toolchain_dir.display());
311            return Ok(libs);
312        }
313    };
314
315    for entry in entries {
316        let entry = entry.into_diagnostic()?;
317        let path = entry.path();
318        let Some(ext) = path.extension() else {
319            continue;
320        };
321
322        if ext == "masp" {
323            log::debug!(target: "flamegraph", "loading library from sysroot: {}", path.display());
324            let bytes = std::fs::read(&path).into_diagnostic()?;
325            let package = miden_mast_package::Package::read_from_bytes(&bytes).map_err(|e| {
326                Report::msg(format!("failed to load package '{}': {e}", path.display()))
327            })?;
328            libs.push(package.mast.clone());
329        } else if ext == "masl" {
330            log::debug!(target: "flamegraph", "loading library from sysroot: {}", path.display());
331            let bytes = std::fs::read(&path).into_diagnostic()?;
332            let lib = Library::read_from_bytes(&bytes).map_err(|e| {
333                Report::msg(format!("failed to load library '{}': {e}", path.display()))
334            })?;
335            libs.push(Arc::new(lib));
336        }
337    }
338
339    Ok(libs)
340}
341
342fn load_package(config: &DebuggerConfig) -> Result<Arc<miden_mast_package::Package>, Report> {
343    let input = config.input.as_ref().ok_or_else(|| Report::msg("no input file specified"))?;
344    let package = match input {
345        InputFile::Real(path) => {
346            let bytes = std::fs::read(path).into_diagnostic()?;
347            miden_mast_package::Package::read_from_bytes(&bytes)
348                .map(Arc::new)
349                .map_err(|e| {
350                    Report::msg(format!(
351                        "failed to load Miden package from {}: {e}",
352                        path.display()
353                    ))
354                })?
355        }
356        InputFile::Stdin(bytes) => miden_mast_package::Package::read_from_bytes(bytes)
357            .map(Arc::new)
358            .map_err(|e| Report::msg(format!("failed to load Miden package from stdin: {e}")))?,
359    };
360
361    if let Some(entry) = config.entrypoint.as_ref() {
362        let id = entry
363            .parse::<miden_assembly::ast::QualifiedProcedureName>()
364            .map_err(|_| Report::msg(format!("invalid function identifier: '{entry}'")))?;
365        if !package.is_library() {
366            return Err(Report::msg("cannot use --entrypoint with executable packages"));
367        }
368
369        package.make_executable(&id).map(Arc::new)
370    } else {
371        Ok(package)
372    }
373}
374
375fn build_stack_path(frames: &[CallFrame]) -> String {
376    build_stack_path_from_names(frames.iter().filter_map(|frame| frame.procedure("")))
377}
378
379fn build_stack_path_from_names<I, S>(frames: I) -> String
380where
381    I: IntoIterator<Item = S>,
382    S: AsRef<str>,
383{
384    let mut path = String::new();
385    for frame in frames {
386        append_frame(&mut path, frame.as_ref());
387    }
388
389    if path.is_empty() {
390        "[unknown]".to_string()
391    } else {
392        path
393    }
394}
395
396fn append_frame(path: &mut String, name: &str) {
397    if !path.is_empty() {
398        path.push(';');
399    }
400    append_sanitized_frame(path, name);
401}
402
403fn append_sanitized_frame(path: &mut String, name: &str) {
404    for ch in name.chars() {
405        match ch {
406            ';' => path.push(':'),
407            '\n' | '\r' => path.push(' '),
408            _ => path.push(ch),
409        }
410    }
411}
412
413fn is_svg_path(path: &Path) -> bool {
414    path.extension()
415        .and_then(|ext| ext.to_str())
416        .is_some_and(|ext| ext.eq_ignore_ascii_case("svg"))
417}
418
419fn write_folded_stacks(samples: &Samples, path: &Path) -> Result<(), Report> {
420    let file = File::create(path).into_diagnostic()?;
421    let mut writer = BufWriter::new(file);
422    for (stack, count) in samples {
423        writeln!(writer, "{stack} {count}").into_diagnostic()?;
424    }
425    writer.flush().into_diagnostic()?;
426
427    eprintln!("Wrote folded stacks to {}", path.display());
428    Ok(())
429}
430
431fn generate_svg(samples: &Samples, path: &Path) -> Result<(), Report> {
432    let input = samples
433        .iter()
434        .map(|(stack, count)| format!("{stack} {count}"))
435        .collect::<Vec<_>>()
436        .join("\n");
437
438    let mut opts = inferno::flamegraph::Options::default();
439    opts.title = "Miden VM Flame Graph (cycles)".to_string();
440    opts.count_name = "cycles".to_string();
441
442    let file = File::create(path).into_diagnostic()?;
443    let mut writer = BufWriter::new(file);
444    inferno::flamegraph::from_reader(&mut opts, input.as_bytes(), &mut writer).into_diagnostic()?;
445    writer.flush().into_diagnostic()?;
446
447    eprintln!("Wrote flame graph to {}", path.display());
448    Ok(())
449}
450
451#[cfg(test)]
452mod tests {
453    use super::FlamegraphProfile;
454
455    #[test]
456    fn record_stack_sanitizes_folded_stack_frames() {
457        let mut profile = FlamegraphProfile::default();
458
459        profile.record_stack(["root;proc", "child\nproc"], 3);
460        profile.record_stack(["root;proc", "child\nproc"], 2);
461        profile.record_stack(["ignored"], 0);
462
463        assert_eq!(profile.total_cycles(), 5);
464        assert_eq!(profile.unique_stack_paths(), 1);
465        assert_eq!(profile.samples().get("root:proc;child proc"), Some(&5));
466        assert!(!profile.samples().contains_key("ignored"));
467    }
468}
miden_debug/flamegraph.rs

miden_debug/
flamegraph.rs
