use std::collections::hash_map;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Write as _;
use std::mem;
use std::ops::Index;
use std::ops::IndexMut;
use crate::global_root::roots;
use crate::key::Key;
use crate::visitor::NodeKind;
use crate::visitor::Visitor;
use crate::visitor::VisitorImpl;
use crate::Allocative;
#[derive(Debug, Default, Clone)]
pub struct FlameGraph {
children: HashMap<Key, FlameGraph>,
children_size: usize,
node_size: usize,
}
impl FlameGraph {
pub fn total_size(&self) -> usize {
self.node_size + self.children_size
}
pub fn add(&mut self, other: FlameGraph) {
self.node_size += other.node_size;
for (key, child) in other.children {
self.add_child(key, child);
}
}
pub fn add_child(&mut self, key: Key, child: FlameGraph) {
self.children_size += child.total_size();
match self.children.entry(key) {
hash_map::Entry::Occupied(mut entry) => {
entry.get_mut().add(child);
}
hash_map::Entry::Vacant(entry) => {
entry.insert(child);
}
}
}
pub fn add_self(&mut self, size: usize) {
self.node_size += size;
}
fn write_flame_graph_impl(&self, stack: &[&str], w: &mut String) {
if self.node_size != 0 {
if !stack.is_empty() {
writeln!(w, "{} {}", stack.join(";"), self.node_size).unwrap();
} else {
}
}
let mut stack = stack.to_vec();
let mut children = Vec::from_iter(self.children.iter());
children.sort_by_key(|(key, _)| *key);
for (key, child) in children {
stack.push(key);
child.write_flame_graph_impl(&stack, w);
stack.pop().unwrap();
}
}
pub fn write(&self) -> String {
let mut r = String::new();
self.write_flame_graph_impl(&[], &mut r);
r
}
}
#[derive(Debug)]
pub struct FlameGraphOutput {
flamegraph: FlameGraph,
warnings: String,
}
impl FlameGraphOutput {
pub fn flamegraph(&self) -> &FlameGraph {
&self.flamegraph
}
pub fn warnings(&self) -> String {
self.warnings.clone()
}
}
#[derive(Default, Eq, PartialEq, Clone, Debug)]
struct TreeData {
size: usize,
rem_size: isize,
unique: bool,
children: HashMap<Key, TreeId>,
}
impl TreeData {
fn inline_children_size(&self) -> isize {
self.size as isize - self.rem_size
}
}
struct TreeRef<'a> {
trees: &'a Trees,
tree_id: TreeId,
}
impl<'a> TreeRef<'a> {
fn write_flame_graph(&self, stack: &[&str], warnings: &mut String) -> FlameGraph {
let mut flame_graph = FlameGraph::default();
let tree = &self.trees[self.tree_id];
if tree.rem_size > 0 {
if stack.is_empty() {
} else {
flame_graph.node_size = tree.rem_size as usize;
}
} else if tree.rem_size < 0 && !stack.is_empty() {
writeln!(
warnings,
"Incorrect size declaration for node `{}`, size of self: {}, size of inline children: {}",
stack.join(";"),
tree.size,
tree.inline_children_size()
)
.unwrap();
}
let mut children: Vec<(&Key, &TreeId)> = Vec::from_iter(&tree.children);
let mut stack = stack.to_vec();
children.sort_by_key(|(k, _)| *k);
for (key, child) in children {
stack.push(key);
let child = TreeRef {
trees: self.trees,
tree_id: *child,
};
let child_framegraph = child.write_flame_graph(&stack, warnings);
flame_graph.add_child(key.clone(), child_framegraph);
stack.pop().unwrap();
}
flame_graph
}
fn to_flame_graph(&self) -> (FlameGraph, String) {
let mut warnings = String::new();
let flame_graph = self.write_flame_graph(&[], &mut warnings);
(flame_graph, warnings)
}
}
#[derive(Debug, Eq, PartialEq)]
struct Tree {
trees: Trees,
tree_id: TreeId,
}
impl Tree {
fn as_ref(&self) -> TreeRef {
TreeRef {
trees: &self.trees,
tree_id: self.tree_id,
}
}
fn to_flame_graph(&self) -> (FlameGraph, String) {
self.as_ref().to_flame_graph()
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
struct TreeId(usize);
#[derive(Default, Clone, Debug, Eq, PartialEq)]
struct Trees {
trees: Vec<TreeData>,
}
impl Trees {
fn new_tree(&mut self) -> TreeId {
let id = TreeId(self.trees.len());
self.trees.push(TreeData::default());
id
}
}
impl Index<TreeId> for Trees {
type Output = TreeData;
fn index(&self, index: TreeId) -> &Self::Output {
&self.trees[index.0]
}
}
impl IndexMut<TreeId> for Trees {
fn index_mut(&mut self, index: TreeId) -> &mut Self::Output {
&mut self.trees[index.0]
}
}
#[derive(Clone, Debug)]
struct TreeStack {
stack: Vec<TreeId>,
tree: TreeId,
}
struct TreeStackRef<'t, 's> {
trees: &'t mut Trees,
stack: &'s mut TreeStack,
}
impl<'t, 's> TreeStackRef<'t, 's> {
fn current_data(&'t mut self) -> &'t mut TreeData {
&mut self.trees[self.stack.tree]
}
fn down(&mut self, key: Key) {
self.stack.stack.push(self.stack.tree);
let next_tree_id = TreeId(self.trees.trees.len());
let child = match self.trees[self.stack.tree].children.entry(key) {
hash_map::Entry::Occupied(e) => *e.get(),
hash_map::Entry::Vacant(e) => {
e.insert(next_tree_id);
let child = self.trees.new_tree();
assert_eq!(child, next_tree_id);
child
}
};
self.stack.tree = child;
}
#[must_use]
fn up(&mut self) -> bool {
if let Some(pop) = self.stack.stack.pop() {
self.stack.tree = pop;
true
} else {
false
}
}
}
#[derive(Eq, PartialEq, Hash, Clone, Copy, Debug)]
struct VisitedSharedPointer(*const ());
unsafe impl Send for VisitedSharedPointer {}
#[derive(Debug)]
pub struct FlameGraphBuilder {
visited_shared: HashSet<VisitedSharedPointer>,
trees: Trees,
current: TreeStack,
shared: Vec<TreeStack>,
root: TreeId,
entered_root_visitor: bool,
}
fn _assert_flame_graph_builder_is_send() {
fn assert_send<T: Send>() {}
assert_send::<FlameGraphBuilder>();
}
impl Default for FlameGraphBuilder {
fn default() -> FlameGraphBuilder {
let mut trees = Trees::default();
let root = trees.new_tree();
FlameGraphBuilder {
trees,
visited_shared: HashSet::new(),
current: TreeStack {
stack: Vec::new(),
tree: root,
},
shared: Vec::new(),
root,
entered_root_visitor: false,
}
}
}
impl FlameGraphBuilder {
pub fn root_visitor(&mut self) -> Visitor {
assert!(!self.entered_root_visitor);
self.entered_root_visitor = true;
Visitor {
visitor: self,
node_kind: NodeKind::Root,
}
}
pub fn visit_root(&mut self, root: &dyn Allocative) {
let mut visitor = self.root_visitor();
root.visit(&mut visitor);
visitor.exit();
}
pub fn visit_global_roots(&mut self) {
for root in roots() {
self.visit_root(root);
}
}
fn finish_impl(mut self) -> Tree {
assert!(self.shared.is_empty());
assert!(self.current.stack.is_empty());
assert!(!self.entered_root_visitor);
Self::update_sizes(self.root, &mut self.trees);
Tree {
trees: self.trees,
tree_id: self.root,
}
}
pub fn finish(self) -> FlameGraphOutput {
let tree = self.finish_impl();
let (flamegraph, warnings) = tree.to_flame_graph();
FlameGraphOutput {
flamegraph,
warnings,
}
}
pub fn finish_and_write_flame_graph(self) -> String {
self.finish().flamegraph.write()
}
fn update_sizes(tree_id: TreeId, trees: &mut Trees) {
let tree = &mut trees[tree_id];
for child in tree.children.values().copied().collect::<Vec<_>>() {
Self::update_sizes(child, trees);
}
let tree = &mut trees[tree_id];
let children_size = if tree.unique {
0
} else {
let tree = &trees[tree_id];
tree.children
.values()
.map(|child| trees[*child].size)
.sum::<usize>()
};
let tree = &mut trees[tree_id];
let size = tree.size;
tree.rem_size = (size as isize).saturating_sub(children_size as isize);
}
fn current(&mut self) -> TreeStackRef {
TreeStackRef {
trees: &mut self.trees,
stack: &mut self.current,
}
}
fn exit_impl(&mut self) {
assert!(self.entered_root_visitor);
let up = self.current().up();
if !up {
if let Some(shared) = self.shared.pop() {
self.current = shared;
assert!(self.current().up());
} else {
self.entered_root_visitor = false;
}
}
}
}
impl VisitorImpl for FlameGraphBuilder {
fn enter_inline_impl(&mut self, name: Key, size: usize, _parent: NodeKind) {
self.current().down(name);
self.current().current_data().size += size;
}
fn enter_unique_impl(&mut self, name: Key, size: usize, _parent: NodeKind) {
self.current().down(name);
self.current().current_data().size += size;
self.current().current_data().unique = true;
}
#[must_use]
fn enter_shared_impl(
&mut self,
name: Key,
size: usize,
ptr: *const (),
_parent: NodeKind,
) -> bool {
self.current().down(name);
self.current().current_data().size += size;
if !self.visited_shared.insert(VisitedSharedPointer(ptr)) {
self.exit_impl();
return false;
}
self.shared.push(mem::replace(
&mut self.current,
TreeStack {
stack: Vec::new(),
tree: self.root,
},
));
true
}
fn exit_inline_impl(&mut self) {
self.exit_impl();
}
fn exit_unique_impl(&mut self) {
self.exit_impl();
}
fn exit_shared_impl(&mut self) {
self.exit_impl();
}
fn exit_root_impl(&mut self) {
self.exit_impl();
}
}
#[cfg(test)]
mod tests {
use crate::flamegraph::FlameGraphBuilder;
use crate::flamegraph::Tree;
use crate::flamegraph::Trees;
use crate::key::Key;
use crate::FlameGraph;
#[test]
fn test_empty() {
let mut fg = FlameGraphBuilder::default();
fg.root_visitor().exit();
let tree = fg.finish_impl();
let mut expected_trees = Trees::default();
let expected_id = expected_trees.new_tree();
let expected = Tree {
trees: expected_trees,
tree_id: expected_id,
};
assert_eq!(expected, tree);
assert_eq!("", tree.to_flame_graph().0.write());
}
#[test]
fn test_simple() {
let mut fg = FlameGraphBuilder::default();
fg.root_visitor().visit_simple(Key::new("a"), 10);
let tree = fg.finish_impl();
let mut expected = Trees::default();
let expected_root = expected.new_tree();
let expected_child = expected.new_tree();
expected[expected_root].size = 0;
expected[expected_root].rem_size = -10;
expected[expected_root]
.children
.insert(Key::new("a"), expected_child);
expected[expected_child].size = 10;
expected[expected_child].rem_size = 10;
let expected = Tree {
trees: expected,
tree_id: expected_root,
};
assert_eq!(expected, tree);
assert_eq!("a 10\n", tree.to_flame_graph().0.write());
}
#[test]
fn test_unique() {
let mut fg = FlameGraphBuilder::default();
let mut visitor = fg.root_visitor();
let mut s = visitor.enter(Key::new("Struct"), 10);
s.visit_simple(Key::new("a"), 3);
let mut un = s.enter_unique(Key::new("p"), 6);
un.visit_simple(Key::new("x"), 13);
un.exit();
s.exit();
visitor.exit();
let tree = fg.finish_impl();
assert_eq!(
"\
Struct 1\n\
Struct;a 3\n\
Struct;p 6\n\
Struct;p;x 13\n\
",
tree.to_flame_graph().0.write(),
"{:#?}",
tree,
);
}
#[test]
fn test_shared() {
let p = 10;
let mut fg = FlameGraphBuilder::default();
let mut visitor = fg.root_visitor();
for _ in 0..2 {
let mut s = visitor.enter(Key::new("Struct"), 10);
s.visit_simple(Key::new("a"), 3);
{
let sh = s.enter_shared(Key::new("p"), 6, &p as *const i32 as *const ());
if let Some(mut sh) = sh {
sh.visit_simple(Key::new("Shared"), 13);
sh.exit();
}
}
s.exit();
}
visitor.exit();
let tree = fg.finish_impl();
assert_eq!(
"\
Shared 13\n\
Struct 2\n\
Struct;a 6\n\
Struct;p 12\n\
",
tree.to_flame_graph().0.write(),
"{:#?}",
tree,
);
}
#[test]
fn test_inline_children_too_large() {
let mut fg = FlameGraphBuilder::default();
let mut visitor = fg.root_visitor();
let mut child_visitor = visitor.enter(Key::new("a"), 10);
child_visitor.visit_simple(Key::new("b"), 13);
child_visitor.exit();
visitor.exit();
let output = fg.finish();
assert_eq!("a;b 13\n", output.flamegraph().write());
assert_eq!(
"Incorrect size declaration for node `a`, size of self: 10, size of inline children: 13\n",
output.warnings()
);
}
#[test]
fn test_flamegraph_add() {
let mut a = FlameGraph::default();
let mut b_1 = FlameGraph::default();
b_1.add_self(10);
let mut b = FlameGraph::default();
b.add_child(Key::new("x"), b_1);
a.add(b);
assert_eq!(10, a.total_size());
}
}