Skip to main content

turbo_trace_server/
bottom_up.rs

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