stak_profiler/
flamegraph.rs

1use crate::{DurationRecord, Error, Stack, StackedRecord};
2use alloc::collections::BTreeMap;
3use std::io::Write;
4
5/// Calculates a flamegraph.
6pub fn calculate_flamegraph(
7    records: impl IntoIterator<Item = Result<DurationRecord, Error>>,
8    mut writer: impl Write,
9) -> Result<(), Error> {
10    let mut map = BTreeMap::<Stack, u128>::new();
11
12    for record in records {
13        let record = record?;
14
15        if let Some(time) = map.get_mut(record.stack()) {
16            *time += record.time();
17        } else {
18            map.insert(record.stack().clone(), record.time());
19        }
20    }
21
22    for (stack, time) in map {
23        writeln!(writer, "{} {}", stack.display_local("<local>"), time)?;
24    }
25
26    Ok(())
27}
28
29#[cfg(test)]
30mod tests {
31    use super::*;
32    use indoc::indoc;
33    use pretty_assertions::assert_eq;
34
35    #[test]
36    fn test_calculate_flamegraph() {
37        let mut buffer = vec![];
38
39        calculate_flamegraph(
40            [
41                Ok(DurationRecord::new(
42                    Stack::new(vec![None, Some("foo".into()), Some("bar".into())]),
43                    1,
44                )),
45                Ok(DurationRecord::new(
46                    Stack::new(vec![None, Some("foo".into()), Some("baz".into())]),
47                    2,
48                )),
49                Ok(DurationRecord::new(
50                    Stack::new(vec![None, Some("foo".into())]),
51                    3,
52                )),
53                Ok(DurationRecord::new(
54                    Stack::new(vec![None, Some("qux".into())]),
55                    4,
56                )),
57                Ok(DurationRecord::new(
58                    Stack::new(vec![None, Some("foo".into()), Some("baz".into())]),
59                    42,
60                )),
61            ],
62            &mut buffer,
63        )
64        .unwrap();
65
66        assert_eq!(
67            String::from_utf8(buffer).unwrap(),
68            indoc!(
69                "
70                <local>;foo 3
71                <local>;foo;bar 1
72                <local>;foo;baz 44
73                <local>;qux 4
74                "
75            )
76        );
77    }
78}
stak_profiler/flamegraph.rs

stak_profiler/
flamegraph.rs
