iai_callgrind_runner/runner/callgrind/
flamegraph.rs

1use std::borrow::Cow;
2use std::cmp::Ordering;
3use std::fs::File;
4use std::io::{BufWriter, Cursor, Write as IoWrite};
5use std::path::{Path, PathBuf};
6
7use anyhow::{Context, Result};
8use inferno::flamegraph::{Direction, Options};
9
10use super::flamegraph_parser::{FlamegraphMap, FlamegraphParser};
11use super::parser::{CallgrindParser, CallgrindProperties, Sentinel};
12use crate::api::{self, EventKind, FlamegraphKind};
13use crate::runner::summary::{BaselineKind, BaselineName, FlamegraphSummaries, FlamegraphSummary};
14use crate::runner::tool::{ToolOutputPath, ToolOutputPathKind};
15
16type ParserOutput = Vec<(PathBuf, CallgrindProperties, FlamegraphMap)>;
17
18#[derive(Debug)]
19pub struct BaselineFlamegraphGenerator {
20    pub baseline_kind: BaselineKind,
21}
22
23#[derive(Debug, Clone)]
24#[allow(clippy::struct_excessive_bools)]
25pub struct Config {
26    pub kind: FlamegraphKind,
27    pub negate_differential: bool,
28    pub normalize_differential: bool,
29    pub event_kinds: Vec<EventKind>,
30    pub direction: Direction,
31    pub title: Option<String>,
32    pub subtitle: Option<String>,
33    pub min_width: f64,
34}
35
36#[derive(Debug, Clone)]
37pub struct Flamegraph {
38    pub config: Config,
39}
40
41#[derive(Debug)]
42pub struct LoadBaselineFlamegraphGenerator {
43    pub loaded_baseline: BaselineName,
44    pub baseline: BaselineName,
45}
46
47#[derive(Debug, Clone)]
48struct OutputPath {
49    pub kind: OutputPathKind,
50    pub event_kind: EventKind,
51    pub baseline_kind: BaselineKind,
52    pub dir: PathBuf,
53    pub name: String,
54    pub modifiers: Vec<String>,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq)]
58enum OutputPathKind {
59    Regular,
60    Old,
61    Base(String),
62    DiffOld,
63    DiffBase(String),
64    DiffBases(String, String),
65}
66
67#[derive(Debug)]
68pub struct SaveBaselineFlamegraphGenerator {
69    pub baseline: BaselineName,
70}
71
72pub trait FlamegraphGenerator {
73    fn create(
74        &self,
75        flamegraph: &Flamegraph,
76        tool_output_path: &ToolOutputPath,
77        sentinel: Option<&Sentinel>,
78        project_root: &Path,
79    ) -> Result<Vec<FlamegraphSummary>>;
80}
81
82impl From<api::FlamegraphConfig> for Config {
83    fn from(value: api::FlamegraphConfig) -> Self {
84        Self {
85            kind: value.kind.unwrap_or(FlamegraphKind::All),
86            negate_differential: value.negate_differential.unwrap_or_default(),
87            normalize_differential: value.normalize_differential.unwrap_or(false),
88            event_kinds: value.event_kinds.unwrap_or_else(|| vec![EventKind::Ir]),
89            direction: value
90                .direction
91                .map_or_else(|| Direction::Inverted, std::convert::Into::into),
92            title: value.title.clone(),
93            subtitle: value.subtitle.clone(),
94            min_width: value.min_width.unwrap_or(0.1f64),
95        }
96    }
97}
98
99impl From<api::Direction> for Direction {
100    fn from(value: api::Direction) -> Self {
101        match value {
102            api::Direction::TopToBottom => Direction::Inverted,
103            api::Direction::BottomToTop => Direction::Straight,
104        }
105    }
106}
107
108impl FlamegraphGenerator for BaselineFlamegraphGenerator {
109    fn create(
110        &self,
111        flamegraph: &Flamegraph,
112        tool_output_path: &ToolOutputPath,
113        sentinel: Option<&Sentinel>,
114        project_root: &Path,
115    ) -> Result<Vec<FlamegraphSummary>> {
116        // We need the dummy path just to clean up and organize the output files independently of
117        // the EventKind of the OutputPath
118        let mut output_path = OutputPath::new(tool_output_path, EventKind::Ir);
119        output_path.init()?;
120        output_path.to_diff_path().clear(true)?;
121        output_path.shift(true)?;
122        output_path.set_modifiers(["total"]);
123
124        if flamegraph.config.kind == FlamegraphKind::None
125            || flamegraph.config.event_kinds.is_empty()
126        {
127            return Ok(vec![]);
128        }
129
130        let (maps, base_maps) =
131            flamegraph.parse(tool_output_path, sentinel, project_root, false)?;
132
133        let total = total_flamegraph_map_from_parsed(&maps).unwrap();
134
135        let mut flamegraph_summaries = FlamegraphSummaries::default();
136        for event_kind in &flamegraph.config.event_kinds {
137            let mut flamegraph_summary = FlamegraphSummary::new(*event_kind);
138            output_path.set_event_kind(*event_kind);
139
140            let stacks_lines = total.to_stack_format(event_kind)?;
141            if flamegraph.is_regular() {
142                Flamegraph::write(
143                    &output_path,
144                    &mut flamegraph.options(*event_kind, output_path.file_name()),
145                    stacks_lines.iter().map(std::string::String::as_str),
146                )?;
147                flamegraph_summary.regular_path = Some(output_path.to_path());
148            }
149
150            if let Some(base_maps) = &base_maps {
151                let total_base = total_flamegraph_map_from_parsed(base_maps).unwrap();
152                // Is Some if FlamegraphKind::Differential or FlamegraphKind::All
153                Flamegraph::create_differential(
154                    &output_path,
155                    &mut flamegraph.options(*event_kind, output_path.to_diff_path().file_name()),
156                    &total_base,
157                    // This unwrap is safe since we always have differential options if the
158                    // flamegraph kind is differential
159                    flamegraph.differential_options().unwrap(),
160                    *event_kind,
161                    &stacks_lines,
162                )?;
163
164                flamegraph_summary.base_path = Some(output_path.to_base_path().to_path());
165                flamegraph_summary.diff_path = Some(output_path.to_diff_path().to_path());
166            }
167
168            flamegraph_summaries.totals.push(flamegraph_summary);
169        }
170
171        Ok(flamegraph_summaries.totals)
172    }
173}
174
175impl Flamegraph {
176    pub fn new(heading: String, mut config: Config) -> Self {
177        if config.title.is_none() {
178            config.title = Some(heading);
179        }
180
181        Self { config }
182    }
183
184    pub fn is_differential(&self) -> bool {
185        matches!(
186            self.config.kind,
187            FlamegraphKind::Differential | FlamegraphKind::All
188        )
189    }
190
191    pub fn is_regular(&self) -> bool {
192        matches!(
193            self.config.kind,
194            FlamegraphKind::Regular | FlamegraphKind::All
195        )
196    }
197
198    pub fn options(&self, event_kind: EventKind, subtitle: String) -> Options {
199        let mut options = Options::default();
200        options.negate_differentials = self.config.negate_differential;
201        options.direction = self.config.direction;
202        options.title.clone_from(
203            self.config
204                .title
205                .as_ref()
206                .expect("A title must be present at this point"),
207        );
208
209        options.subtitle = if let Some(subtitle) = &self.config.subtitle {
210            Some(subtitle.clone())
211        } else {
212            Some(subtitle)
213        };
214
215        options.min_width = self.config.min_width;
216        options.count_name = event_kind.to_string();
217        options
218    }
219
220    pub fn differential_options(&self) -> Option<inferno::differential::Options> {
221        self.is_differential()
222            .then(|| inferno::differential::Options {
223                normalize: self.config.normalize_differential,
224                ..Default::default()
225            })
226    }
227
228    pub fn parse<P>(
229        &self,
230        tool_output_path: &ToolOutputPath,
231        sentinel: Option<&Sentinel>,
232        project_root: P,
233        no_differential: bool,
234    ) -> Result<(ParserOutput, Option<ParserOutput>)>
235    where
236        P: Into<PathBuf>,
237    {
238        let parser = FlamegraphParser::new(sentinel, project_root);
239        // We need this map in all remaining cases of `FlamegraphKinds`
240        let mut maps = parser.parse(tool_output_path)?;
241
242        let base_path = tool_output_path.to_base_path();
243        let mut base_maps = (!no_differential && self.is_differential() && base_path.exists())
244            .then(|| parser.parse(&base_path))
245            .transpose()?;
246
247        if self.config.event_kinds.iter().any(EventKind::is_derived) {
248            for map in &mut maps {
249                map.2.make_summary()?;
250            }
251            if let Some(maps) = base_maps.as_mut() {
252                for map in maps {
253                    map.2.make_summary()?;
254                }
255            }
256        }
257
258        Ok((maps, base_maps))
259    }
260
261    fn create_differential(
262        output_path: &OutputPath,
263        options: &mut inferno::flamegraph::Options,
264        base_map: &FlamegraphMap,
265        differential_options: inferno::differential::Options,
266        event_kind: EventKind,
267        stacks_lines: &[String],
268    ) -> Result<()> {
269        let base_stacks_lines = base_map.to_stack_format(&event_kind)?;
270
271        let cursor = Cursor::new(stacks_lines.join("\n"));
272        let base_cursor = Cursor::new(base_stacks_lines.join("\n"));
273        let mut result = Cursor::new(vec![]);
274
275        inferno::differential::from_readers(differential_options, base_cursor, cursor, &mut result)
276            .context("Failed creating a differential flamegraph")?;
277
278        let diff_output_path = output_path.to_diff_path();
279        Flamegraph::write(
280            &diff_output_path,
281            options,
282            String::from_utf8_lossy(result.get_ref()).lines(),
283        )
284    }
285
286    fn write<'stacks>(
287        output_path: &OutputPath,
288        options: &mut Options<'_>,
289        stacks: impl Iterator<Item = &'stacks str>,
290    ) -> Result<()> {
291        let path = output_path.to_path();
292        let mut writer = BufWriter::new(output_path.create()?);
293        inferno::flamegraph::from_lines(options, stacks, &mut writer)
294            .with_context(|| format!("Failed creating a flamegraph at '{}'", path.display()))?;
295
296        writer
297            .flush()
298            .with_context(|| format!("Failed flushing content to '{}'", path.display()))
299    }
300}
301
302impl FlamegraphGenerator for LoadBaselineFlamegraphGenerator {
303    fn create(
304        &self,
305        flamegraph: &Flamegraph,
306        tool_output_path: &ToolOutputPath,
307        sentinel: Option<&Sentinel>,
308        project_root: &Path,
309    ) -> Result<Vec<FlamegraphSummary>> {
310        // We need the dummy path just to clean up and organize the output files independently of
311        // the EventKind of the OutputPath
312        let mut output_path = OutputPath::new(tool_output_path, EventKind::Ir);
313
314        if flamegraph.config.kind == FlamegraphKind::None
315            || flamegraph.config.event_kinds.is_empty()
316            || !flamegraph.is_differential()
317        {
318            return Ok(vec![]);
319        }
320
321        output_path.to_diff_path().clear(true)?;
322        output_path.set_modifiers(["total"]);
323
324        let (maps, base_maps) = flamegraph
325            .parse(tool_output_path, sentinel, project_root, false)
326            .map(|(a, b)| (a, b.unwrap()))?;
327
328        let mut flamegraph_summaries = FlamegraphSummaries::default();
329        if let Some(total) = total_flamegraph_map_from_parsed(&maps) {
330            let base_total = total_flamegraph_map_from_parsed(&base_maps);
331
332            if let Some(base_total) = base_total {
333                for event_kind in &flamegraph.config.event_kinds {
334                    let mut flamegraph_summary = FlamegraphSummary::new(*event_kind);
335                    output_path.set_event_kind(*event_kind);
336
337                    Flamegraph::create_differential(
338                        &output_path,
339                        &mut flamegraph
340                            .options(*event_kind, output_path.to_diff_path().file_name()),
341                        &base_total,
342                        // This unwrap is safe since we always produce a differential flamegraph
343                        flamegraph.differential_options().unwrap(),
344                        *event_kind,
345                        &total.to_stack_format(event_kind)?,
346                    )?;
347
348                    flamegraph_summary.regular_path = Some(output_path.to_path());
349                    flamegraph_summary.base_path = Some(output_path.to_base_path().to_path());
350                    flamegraph_summary.diff_path = Some(output_path.to_diff_path().to_path());
351
352                    flamegraph_summaries.totals.push(flamegraph_summary);
353                }
354            }
355        }
356
357        Ok(flamegraph_summaries.totals)
358    }
359}
360
361impl OutputPath {
362    pub fn new(tool_output_path: &ToolOutputPath, event_kind: EventKind) -> Self {
363        Self {
364            kind: match &tool_output_path.kind {
365                ToolOutputPathKind::Out | ToolOutputPathKind::Log => OutputPathKind::Regular,
366                ToolOutputPathKind::OldOut | ToolOutputPathKind::OldLog => OutputPathKind::Old,
367                ToolOutputPathKind::BaseLog(name) | ToolOutputPathKind::Base(name) => {
368                    OutputPathKind::Base(name.clone())
369                }
370            },
371            event_kind,
372            baseline_kind: tool_output_path.baseline_kind.clone(),
373            dir: tool_output_path.dir.clone(),
374            name: tool_output_path.name.clone(),
375            modifiers: Vec::default(),
376        }
377    }
378
379    pub fn init(&self) -> Result<()> {
380        std::fs::create_dir_all(&self.dir)
381            .with_context(|| {
382                format!(
383                    "Failed creating flamegraph directory '{}'",
384                    self.dir.display()
385                )
386            })
387            .map_err(Into::into)
388    }
389
390    pub fn create(&self) -> Result<File> {
391        let path = self.to_path();
392        File::create(&path)
393            .with_context(|| format!("Failed creating flamegraph file '{}'", path.display()))
394    }
395
396    pub fn clear(&self, ignore_event_kind: bool) -> Result<()> {
397        for path in self.real_paths(ignore_event_kind)? {
398            std::fs::remove_file(path)?;
399        }
400
401        Ok(())
402    }
403
404    /// This method will remove all differential flamegraphs for a specific base or old
405    ///
406    /// The differential flamegraphs with a base can end with the base name
407    /// (`*.diff.base@<name>.svg`) and/or with the parts until `flamegraph` removed start with the
408    /// base name (`base@<name>.diff.*`)
409    pub fn clear_diff(&self) -> Result<()> {
410        let extension = match &self.baseline_kind {
411            BaselineKind::Old => "diff.old.svg".to_owned(),
412            BaselineKind::Name(name) => format!("diff.base@{name}.svg"),
413        };
414        for entry in std::fs::read_dir(&self.dir)
415            .with_context(|| format!("Failed reading directory '{}'", self.dir.display()))?
416        {
417            let entry = entry?;
418            let file_name = entry.file_name().to_string_lossy().to_string();
419            if let Some(suffix) =
420                file_name.strip_prefix(format!("callgrind.{}", &self.name).as_str())
421            {
422                let path = entry.path();
423
424                if suffix.ends_with(extension.as_str()) {
425                    std::fs::remove_file(&path).with_context(|| {
426                        format!("Failed removing flamegraph file: '{}'", path.display())
427                    })?;
428                }
429
430                if let BaselineKind::Name(name) = &self.baseline_kind {
431                    if suffix
432                        .split('.')
433                        .skip_while(|p| *p != "flamegraph")
434                        .take(3)
435                        .eq([
436                            "flamegraph".to_owned(),
437                            format!("base@{name}"),
438                            "diff".to_owned(),
439                        ])
440                    {
441                        std::fs::remove_file(&path).with_context(|| {
442                            format!("Failed removing flamegraph file: '{}'", path.display())
443                        })?;
444                    }
445                } else {
446                    // do nothing
447                }
448            }
449        }
450
451        Ok(())
452    }
453
454    pub fn shift(&self, ignore_event_kind: bool) -> Result<()> {
455        match &self.baseline_kind {
456            BaselineKind::Old => {
457                self.to_base_path().clear(ignore_event_kind)?;
458                for path in self.real_paths(ignore_event_kind)? {
459                    let new_path = path.with_extension("old.svg");
460                    std::fs::rename(&path, &new_path).with_context(|| {
461                        format!(
462                            "Failed moving flamegraph file from '{}' to '{}'",
463                            path.display(),
464                            new_path.display()
465                        )
466                    })?;
467                }
468                Ok(())
469            }
470            BaselineKind::Name(_) => self.clear(ignore_event_kind),
471        }
472    }
473
474    pub fn to_diff_path(&self) -> Self {
475        Self {
476            kind: match (&self.kind, &self.baseline_kind) {
477                (OutputPathKind::Regular, BaselineKind::Old) => OutputPathKind::DiffOld,
478                (OutputPathKind::Regular, BaselineKind::Name(name)) => {
479                    OutputPathKind::DiffBase(name.to_string())
480                }
481                (OutputPathKind::Base(name), BaselineKind::Name(other)) => {
482                    OutputPathKind::DiffBases(name.to_string(), other.to_string())
483                }
484                // TODO: NOT UNREACHABLE?
485                (OutputPathKind::Old | OutputPathKind::Base(_), _) => unreachable!(),
486                (value, _) => value.clone(),
487            },
488            ..self.clone()
489        }
490    }
491
492    pub fn to_base_path(&self) -> Self {
493        Self {
494            kind: match &self.baseline_kind {
495                BaselineKind::Old => OutputPathKind::Old,
496                BaselineKind::Name(name) => OutputPathKind::Base(name.to_string()),
497            },
498            ..self.clone()
499        }
500    }
501
502    pub fn extension(&self) -> String {
503        match &self.kind {
504            OutputPathKind::Regular => format!("{}.flamegraph.svg", self.event_kind.to_name()),
505            OutputPathKind::Old => format!("{}.flamegraph.old.svg", self.event_kind.to_name()),
506            OutputPathKind::Base(name) => {
507                format!("{}.flamegraph.base@{name}.svg", self.event_kind.to_name())
508            }
509            OutputPathKind::DiffOld => {
510                format!("{}.flamegraph.diff.old.svg", self.event_kind.to_name())
511            }
512            OutputPathKind::DiffBase(name) => {
513                format!(
514                    "{}.flamegraph.diff.base@{name}.svg",
515                    self.event_kind.to_name()
516                )
517            }
518            OutputPathKind::DiffBases(name, base) => {
519                format!(
520                    "{}.flamegraph.base@{name}.diff.base@{base}.svg",
521                    self.event_kind.to_name()
522                )
523            }
524        }
525    }
526
527    pub fn set_modifiers<I, T>(&mut self, modifiers: T)
528    where
529        T: IntoIterator<Item = I>,
530        I: Into<String>,
531    {
532        self.modifiers = modifiers.into_iter().map(Into::into).collect();
533    }
534
535    pub fn set_event_kind(&mut self, event_kind: EventKind) {
536        self.event_kind = event_kind;
537    }
538
539    pub fn real_paths(&self, ignore_event_kind: bool) -> Result<Vec<PathBuf>> {
540        let extension = self.extension();
541        let to_match = if ignore_event_kind {
542            extension
543                .split_once('.')
544                .expect("The '.' delimiter should be present at least once")
545                .1
546        } else {
547            &extension
548        };
549
550        let mut paths = vec![];
551        for entry in std::fs::read_dir(&self.dir)
552            .with_context(|| format!("Failed reading directory '{}'", self.dir.display()))?
553        {
554            let path = entry?;
555            let file_name = path.file_name().to_string_lossy().to_string();
556            if let Some(suffix) =
557                file_name.strip_prefix(format!("callgrind.{}.", &self.name).as_str())
558            {
559                if suffix.ends_with(to_match) {
560                    paths.push(path.path());
561                }
562            }
563        }
564
565        Ok(paths)
566    }
567
568    pub fn file_name(&self) -> String {
569        if self.modifiers.is_empty() {
570            format!("callgrind.{}.{}", self.name, self.extension())
571        } else {
572            format!(
573                "callgrind.{}.{}.{}",
574                self.name,
575                self.modifiers.join("."),
576                self.extension()
577            )
578        }
579    }
580
581    pub fn to_path(&self) -> PathBuf {
582        self.dir.join(self.file_name())
583    }
584}
585
586impl FlamegraphGenerator for SaveBaselineFlamegraphGenerator {
587    fn create(
588        &self,
589        flamegraph: &Flamegraph,
590        tool_output_path: &ToolOutputPath,
591        sentinel: Option<&Sentinel>,
592        project_root: &Path,
593    ) -> Result<Vec<FlamegraphSummary>> {
594        // We need the dummy path just to clean up and organize the output files independently of
595        // the EventKind of the OutputPath
596        let mut output_path = OutputPath::new(tool_output_path, EventKind::Ir);
597        output_path.init()?;
598        output_path.clear(true)?;
599        output_path.clear_diff()?;
600        output_path.set_modifiers(["total"]);
601
602        if flamegraph.config.kind == FlamegraphKind::None
603            || flamegraph.config.event_kinds.is_empty()
604            || !flamegraph.is_regular()
605        {
606            return Ok(vec![]);
607        }
608
609        let (maps, _) = flamegraph.parse(tool_output_path, sentinel, project_root, true)?;
610        let total_map = total_flamegraph_map_from_parsed(&maps).unwrap();
611
612        let mut flamegraph_summaries = FlamegraphSummaries::default();
613        for event_kind in &flamegraph.config.event_kinds {
614            let mut flamegraph_summary = FlamegraphSummary::new(*event_kind);
615            output_path.set_event_kind(*event_kind);
616
617            Flamegraph::write(
618                &output_path,
619                &mut flamegraph.options(*event_kind, output_path.file_name()),
620                total_map
621                    .to_stack_format(event_kind)?
622                    .iter()
623                    .map(String::as_str),
624            )?;
625
626            flamegraph_summary.regular_path = Some(output_path.to_path());
627            flamegraph_summaries.summaries.push(flamegraph_summary);
628        }
629
630        Ok(flamegraph_summaries.totals)
631    }
632}
633
634fn total_flamegraph_map_from_parsed(maps: &ParserOutput) -> Option<Cow<FlamegraphMap>> {
635    match maps.len().cmp(&1) {
636        Ordering::Less => None,
637        Ordering::Equal => Some(Cow::Borrowed(&maps[0].2)),
638        Ordering::Greater => {
639            let mut total = maps[0].2.clone();
640            for (_, _, map) in maps.iter().skip(1) {
641                total.add(map);
642            }
643            Some(Cow::Owned(total))
644        }
645    }
646}