turbo_trace_server/
bottom_up.rs

1use std::{env, sync::Arc};
2
3use hashbrown::HashMap;
4
5use crate::{
6    span::{SpanBottomUp, SpanIndex},
7    span_ref::SpanRef,
8    string_tuple_ref::StringTupleRef,
9};
10
11pub struct SpanBottomUpBuilder {
12    // These values won't change after creation:
13    pub self_spans: Vec<SpanIndex>,
14    pub children: HashMap<(String, String), SpanBottomUpBuilder>,
15    pub example_span: SpanIndex,
16}
17
18impl SpanBottomUpBuilder {
19    pub fn new(example_span: SpanIndex) -> Self {
20        Self {
21            self_spans: vec![],
22            children: HashMap::default(),
23            example_span,
24        }
25    }
26
27    pub fn build(self) -> SpanBottomUp {
28        SpanBottomUp::new(
29            self.self_spans,
30            self.example_span,
31            self.children
32                .into_values()
33                .map(|child| Arc::new(child.build()))
34                .collect(),
35        )
36    }
37}
38
39pub fn build_bottom_up_graph<'a>(
40    spans: impl Iterator<Item = SpanRef<'a>>,
41) -> Vec<Arc<SpanBottomUp>> {
42    let max_depth = env::var("BOTTOM_UP_DEPTH")
43        .ok()
44        .and_then(|s| s.parse().ok())
45        .unwrap_or(usize::MAX);
46    let mut roots: HashMap<(String, String), SpanBottomUpBuilder> = HashMap::default();
47
48    // unfortunately there is a rustc bug that fails the typechecking here
49    // when using Either<impl Iterator, impl Iterator>. This error appears
50    // in certain cases when building next-swc.
51    //
52    // see here: https://github.com/rust-lang/rust/issues/124891
53    let mut current_iterators: Vec<Box<dyn Iterator<Item = SpanRef<'_>>>> =
54        vec![Box::new(spans.flat_map(|span| span.children()))];
55
56    let mut current_path: Vec<((&'_ str, &'_ str), SpanIndex)> = vec![];
57    while let Some(mut iter) = current_iterators.pop() {
58        if let Some(child) = iter.next() {
59            current_iterators.push(iter);
60
61            let (category, name) = child.group_name();
62            let (_, mut bottom_up) = roots
63                .raw_entry_mut()
64                .from_key(&StringTupleRef(category, name))
65                .or_insert_with(|| {
66                    (
67                        (category.to_string(), name.to_string()),
68                        SpanBottomUpBuilder::new(child.index()),
69                    )
70                });
71            bottom_up.self_spans.push(child.index());
72            let mut prev = None;
73            for &((category, title), example_span) in current_path.iter().rev().take(max_depth) {
74                if prev == Some((category, title)) {
75                    continue;
76                }
77                let (_, child_bottom_up) = bottom_up
78                    .children
79                    .raw_entry_mut()
80                    .from_key(&StringTupleRef(category, title))
81                    .or_insert_with(|| {
82                        (
83                            (category.to_string(), title.to_string()),
84                            SpanBottomUpBuilder::new(example_span),
85                        )
86                    });
87                child_bottom_up.self_spans.push(child.index());
88                bottom_up = child_bottom_up;
89                prev = Some((category, title));
90            }
91
92            current_path.push((child.group_name(), child.index()));
93            current_iterators.push(Box::new(child.children()));
94        } else {
95            current_path.pop();
96        }
97    }
98    roots.into_values().map(|b| Arc::new(b.build())).collect()
99}