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