miden-debug 0.10.3

use std::{
    collections::BTreeMap,
    fs::File,
    io::{BufWriter, Write},
    path::{Path, PathBuf},
    sync::Arc,
};

use miden_assembly::{DefaultSourceManager, SourceManager};
use miden_assembly_syntax::diagnostics::{IntoDiagnostic, Report, WrapErr};
use miden_debug_engine::LinkLibrary;
use miden_processor::ExecutionError;

use crate::{
    config::{ColorChoice, DebuggerConfig},
    debug::CallFrame,
    exec::{DebugExecutor, ExecutionConfig},
    felt::Felt,
    input::InputFile,
};

/// Folded stack samples keyed by semicolon-separated stack paths.
pub type Samples = BTreeMap<String, usize>;

/// A collected VM cycle profile that can be rendered as folded stacks or an SVG flamegraph.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct FlamegraphProfile {
    samples: Samples,
    total_cycles: usize,
}

/// The output format selected by [`FlamegraphProfile::write_to_path`].
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum FlamegraphOutput {
    Svg,
    FoldedStacks,
}

impl FlamegraphProfile {
    /// Execute `executor` to completion, sampling its call stack for every VM cycle.
    pub fn collect(executor: &mut DebugExecutor) -> Result<Self, ExecutionError> {
        let mut profile = Self::default();

        loop {
            if executor.stopped {
                break;
            }

            let previous_cycle = executor.cycle;
            match executor.step() {
                Ok(_) if executor.cycle > previous_cycle => {
                    let cycle_delta = executor.cycle - previous_cycle;
                    profile.record_call_stack(executor.callstack.frames(), cycle_delta);
                }
                Ok(_) => {
                    if executor.stopped {
                        break;
                    }
                }
                Err(err) => return Err(err),
            }
        }

        Ok(profile)
    }

    /// Record `cycles` against a call stack from the debugger engine.
    pub fn record_call_stack(&mut self, frames: &[CallFrame], cycles: usize) {
        let path = build_stack_path(frames);
        self.record_stack_path(path, cycles);
    }

    /// Record `cycles` against a stack of frame names.
    ///
    /// Frame names are sanitized for folded stack output, then joined with `;`.
    pub fn record_stack<I, S>(&mut self, frames: I, cycles: usize)
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        let path = build_stack_path_from_names(frames);
        self.record_stack_path(path, cycles);
    }

    /// Record `cycles` against an already formatted folded-stack path.
    pub fn record_stack_path(&mut self, stack_path: impl Into<String>, cycles: usize) {
        if cycles == 0 {
            return;
        }

        self.total_cycles += cycles;
        *self.samples.entry(stack_path.into()).or_default() += cycles;
    }

    pub fn samples(&self) -> &Samples {
        &self.samples
    }

    pub fn total_cycles(&self) -> usize {
        self.total_cycles
    }

    pub fn unique_stack_paths(&self) -> usize {
        self.samples.len()
    }

    /// Write this profile as an SVG when `path` ends in `.svg`, otherwise as folded stack text.
    pub fn write_to_path(&self, path: impl AsRef<Path>) -> Result<FlamegraphOutput, Report> {
        let path = path.as_ref();
        if is_svg_path(path) {
            self.write_svg(path)?;
            Ok(FlamegraphOutput::Svg)
        } else {
            self.write_folded_stacks(path)?;
            Ok(FlamegraphOutput::FoldedStacks)
        }
    }

    /// Write this profile in folded stack format.
    pub fn write_folded_stacks(&self, path: impl AsRef<Path>) -> Result<(), Report> {
        write_folded_stacks(&self.samples, path.as_ref())
    }

    /// Render this profile as an SVG flamegraph.
    pub fn write_svg(&self, path: impl AsRef<Path>) -> Result<(), Report> {
        generate_svg(&self.samples, path.as_ref())
    }
}

#[derive(clap::Args, Debug)]
pub struct FlamegraphArgs {
    /// Specify the path to a Miden program file to execute.
    #[arg(value_name = "FILE")]
    pub input: InputFile,
    /// Write the generated flame graph SVG or folded stack text to this path.
    #[arg(short, long, default_value = "flamegraph.svg")]
    pub output: PathBuf,
    /// Specify the path to a file containing program inputs.
    #[arg(long, value_name = "FILE")]
    pub inputs: Option<ExecutionConfig>,
    /// Arguments to place on the operand stack before calling the program entrypoint.
    #[arg(last(true), value_name = "ARGV")]
    pub args: Vec<Felt>,
    /// The working directory for execution.
    #[arg(long, value_name = "DIR", help_heading = "Execution")]
    pub working_dir: Option<PathBuf>,
    /// The path to the root directory of the current Miden toolchain.
    #[arg(
        long,
        value_name = "DIR",
        env = "MIDEN_SYSROOT",
        help_heading = "Linker"
    )]
    pub sysroot: Option<PathBuf>,
    /// Specify the function to call as the entrypoint for the program.
    #[arg(long, help_heading = "Execution")]
    pub entrypoint: Option<String>,
    /// Specify one or more search paths for link libraries requested via `-l`.
    #[arg(
        long = "search-path",
        short = 'L',
        value_name = "PATH",
        help_heading = "Linker"
    )]
    pub search_path: Vec<PathBuf>,
    /// Link compiled projects to the specified library NAME.
    #[arg(
        long = "link-library",
        short = 'l',
        value_name = "[KIND=]NAME",
        value_delimiter = ',',
        next_line_help(true),
        help_heading = "Linker"
    )]
    pub link_libraries: Vec<LinkLibrary>,
    /// Source path prefixes used by the compiler's `-Zremap-path-prefix` option.
    ///
    /// When debug info stores trimmed source paths, DAP clients may still send
    /// absolute editor paths. These prefixes provide an explicit mapping between
    /// the two forms.
    #[cfg(feature = "dap")]
    #[cfg_attr(
        feature = "tui",
        arg(
            long = "source-path-prefix",
            alias = "trim-path-prefix",
            value_name = "PATH",
            help_heading = "Debugging"
        )
    )]
    pub source_path_prefixes: Vec<PathBuf>,
}

impl FlamegraphArgs {
    fn into_debugger_config(self) -> DebuggerConfig {
        DebuggerConfig {
            input: Some(self.input),
            inputs: self.inputs,
            args: self.args,
            working_dir: self.working_dir,
            sysroot: self.sysroot,
            color: ColorChoice::Auto,
            entrypoint: self.entrypoint,
            #[cfg(feature = "dap")]
            dap_connect: None,
            #[cfg(feature = "dap")]
            start_debug_adapter: None,
            #[cfg(feature = "dap")]
            source_path_prefixes: self.source_path_prefixes,
            search_path: self.search_path,
            link_libraries: self.link_libraries,
            repl: false,
            commands: None,
            replay: None,
            #[cfg(feature = "python")]
            no_user_python_init: true,
            profiler_cli_args: Default::default(), // disable profiling
        }
    }
}

pub fn run(args: FlamegraphArgs) -> Result<(), Report> {
    let output = args.output.clone();
    let mut config = args.into_debugger_config();
    ensure_working_dir(&mut config)?;

    let source_manager: Arc<dyn SourceManager> = Arc::new(DefaultSourceManager::default());
    let mut executor =
        crate::program_loader::load_debug_executor(&config, source_manager, "flamegraph")?;

    let profile = match FlamegraphProfile::collect(&mut executor) {
        Ok(profile) => profile,
        Err(err) => {
            return Err(Report::msg(format!(
                "program execution failed at cycle {}: {err}",
                executor.cycle
            )));
        }
    };

    eprintln!(
        "Executed {} cycles across {} unique stack paths",
        profile.total_cycles(),
        profile.unique_stack_paths()
    );

    profile.write_to_path(&output)?;

    Ok(())
}

fn ensure_working_dir(config: &mut DebuggerConfig) -> Result<(), Report> {
    if config.working_dir.is_none() {
        let cwd = std::env::current_dir()
            .into_diagnostic()
            .wrap_err("could not read current working directory")?;
        config.working_dir = Some(cwd);
    }

    Ok(())
}

fn build_stack_path(frames: &[CallFrame]) -> String {
    build_stack_path_from_names(frames.iter().filter_map(|frame| frame.procedure("")))
}

fn build_stack_path_from_names<I, S>(frames: I) -> String
where
    I: IntoIterator<Item = S>,
    S: AsRef<str>,
{
    let mut path = String::new();
    for frame in frames {
        append_frame(&mut path, frame.as_ref());
    }

    if path.is_empty() {
        "[unknown]".to_string()
    } else {
        path
    }
}

fn append_frame(path: &mut String, name: &str) {
    if !path.is_empty() {
        path.push(';');
    }
    append_sanitized_frame(path, name);
}

fn append_sanitized_frame(path: &mut String, name: &str) {
    for ch in name.chars() {
        match ch {
            ';' => path.push(':'),
            '\n' | '\r' => path.push(' '),
            _ => path.push(ch),
        }
    }
}

fn is_svg_path(path: &Path) -> bool {
    path.extension()
        .and_then(|ext| ext.to_str())
        .is_some_and(|ext| ext.eq_ignore_ascii_case("svg"))
}

fn write_folded_stacks(samples: &Samples, path: &Path) -> Result<(), Report> {
    let file = File::create(path).into_diagnostic()?;
    let mut writer = BufWriter::new(file);
    for (stack, count) in samples {
        writeln!(writer, "{stack} {count}").into_diagnostic()?;
    }
    writer.flush().into_diagnostic()?;

    eprintln!("Wrote folded stacks to {}", path.display());
    Ok(())
}

fn generate_svg(samples: &Samples, path: &Path) -> Result<(), Report> {
    let input = samples
        .iter()
        .map(|(stack, count)| format!("{stack} {count}"))
        .collect::<Vec<_>>()
        .join("\n");

    let mut opts = inferno::flamegraph::Options::default();
    opts.title = "Miden VM Flame Graph (cycles)".to_string();
    opts.count_name = "cycles".to_string();

    let file = File::create(path).into_diagnostic()?;
    let mut writer = BufWriter::new(file);
    inferno::flamegraph::from_reader(&mut opts, input.as_bytes(), &mut writer).into_diagnostic()?;
    writer.flush().into_diagnostic()?;

    eprintln!("Wrote flame graph to {}", path.display());
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::FlamegraphProfile;

    #[test]
    fn record_stack_sanitizes_folded_stack_frames() {
        let mut profile = FlamegraphProfile::default();

        profile.record_stack(["root;proc", "child\nproc"], 3);
        profile.record_stack(["root;proc", "child\nproc"], 2);
        profile.record_stack(["ignored"], 0);

        assert_eq!(profile.total_cycles(), 5);
        assert_eq!(profile.unique_stack_paths(), 1);
        assert_eq!(profile.samples().get("root:proc;child proc"), Some(&5));
        assert!(!profile.samples().contains_key("ignored"));
    }
}