stak_profiler/
flamegraph.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
use crate::{DurationRecord, Error, Stack, StackedRecord};
use alloc::collections::BTreeMap;
use std::io::Write;

/// Calculates a flamegraph.
pub fn calculate_flamegraph(
    records: impl IntoIterator<Item = Result<DurationRecord, Error>>,
    mut writer: impl Write,
) -> Result<(), Error> {
    let mut map = BTreeMap::<Stack, u128>::new();

    for record in records {
        let record = record?;

        if let Some(time) = map.get_mut(record.stack()) {
            *time += record.time();
        } else {
            map.insert(record.stack().clone(), record.time());
        }
    }

    for (stack, time) in map {
        writeln!(writer, "{} {}", stack.display_local("<local>"), time)?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use indoc::indoc;
    use pretty_assertions::assert_eq;

    #[test]
    fn test_calculate_flamegraph() {
        let mut buffer = vec![];

        calculate_flamegraph(
            [
                Ok(DurationRecord::new(
                    Stack::new(vec![None, Some("foo".into()), Some("bar".into())]),
                    1,
                )),
                Ok(DurationRecord::new(
                    Stack::new(vec![None, Some("foo".into()), Some("baz".into())]),
                    2,
                )),
                Ok(DurationRecord::new(
                    Stack::new(vec![None, Some("foo".into())]),
                    3,
                )),
                Ok(DurationRecord::new(
                    Stack::new(vec![None, Some("qux".into())]),
                    4,
                )),
                Ok(DurationRecord::new(
                    Stack::new(vec![None, Some("foo".into()), Some("baz".into())]),
                    42,
                )),
            ],
            &mut buffer,
        )
        .unwrap();

        assert_eq!(
            String::from_utf8(buffer).unwrap(),
            indoc!(
                "
                <local>;foo 3
                <local>;foo;bar 1
                <local>;foo;baz 44
                <local>;qux 4
                "
            )
        );
    }
}