fix bottom up tracing for aggregated spans

crates/turbopack-trace-server/src/bottom_up.rs

@@ -1,13 +1,47 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
+
+use either::Either;
 
 use crate::{
     span::{SpanBottomUp, SpanIndex},
     span_ref::SpanRef,
 };
 
-pub fn build_bottom_up_graph(span: SpanRef<'_>) -> Vec<SpanBottomUp> {
+pub struct SpanBottomUpBuilder {
+    // These values won't change after creation:
+    pub self_spans: Vec<SpanIndex>,
+    pub children: HashMap<String, SpanBottomUpBuilder>,
+    pub example_span: SpanIndex,
+}
+
+impl SpanBottomUpBuilder {
+    pub fn new(example_span: SpanIndex) -> Self {
+        Self {
+            self_spans: vec![],
+            children: HashMap::new(),
+            example_span,
+        }
+    }
+
+    pub fn build(self) -> SpanBottomUp {
+        SpanBottomUp::new(
+            self.self_spans,
+            self.example_span,
+            self.children
+                .into_values()
+                .map(|child| Arc::new(child.build()))
+                .collect(),
+        )
+    }
+}
+
+pub fn build_bottom_up_graph<'a>(
+    spans: impl IntoIterator<Item = SpanRef<'a>>,
+) -> Vec<Arc<SpanBottomUp>> {
     let mut roots = HashMap::new();
-    let mut current_iterators = vec![span.children()];
+    let mut current_iterators = vec![Either::Left(
+        spans.into_iter().flat_map(|span| span.children()),
+    )];
     let mut current_path: Vec<(&'_ str, SpanIndex)> = vec![];
     while let Some(mut iter) = current_iterators.pop() {
         if let Some(child) = iter.next() {
@@ -17,7 +51,7 @@ pub fn build_bottom_up_graph(span: SpanRef<'_>) -> Vec<SpanBottomUp> {
             let (_, mut bottom_up) = roots
                 .raw_entry_mut()
                 .from_key(name)
-                .or_insert_with(|| (name.to_string(), SpanBottomUp::new(child.index())));
+                .or_insert_with(|| (name.to_string(), SpanBottomUpBuilder::new(child.index())));
             bottom_up.self_spans.push(child.index());
             let mut prev = None;
             for &(name, example_span) in current_path.iter().rev() {
@@ -28,17 +62,17 @@ pub fn build_bottom_up_graph(span: SpanRef<'_>) -> Vec<SpanBottomUp> {
                     .children
                     .raw_entry_mut()
                     .from_key(name)
-                    .or_insert_with(|| (name.to_string(), SpanBottomUp::new(example_span)));
+                    .or_insert_with(|| (name.to_string(), SpanBottomUpBuilder::new(example_span)));
                 child_bottom_up.self_spans.push(child.index());
                 bottom_up = child_bottom_up;
                 prev = Some(name);
             }
 
             current_path.push((child.group_name(), child.index()));
-            current_iterators.push(child.children());
+            current_iterators.push(Either::Right(child.children()));
         } else {
             current_path.pop();
         }
     }
-    roots.into_values().collect()
+    roots.into_values().map(|b| Arc::new(b.build())).collect()
 }

crates/turbopack-trace-server/src/span.rs

@@ -41,7 +41,7 @@ pub struct Span {
     pub corrected_self_time: OnceLock<u64>,
     pub corrected_total_time: OnceLock<u64>,
     pub graph: OnceLock<Vec<SpanGraphEvent>>,
-    pub bottom_up: OnceLock<Vec<SpanBottomUp>>,
+    pub bottom_up: OnceLock<Vec<Arc<SpanBottomUp>>>,
     pub search_index: OnceLock<HashMap<String, Vec<SpanIndex>>>,
 }
 
@@ -83,12 +83,13 @@ pub struct SpanGraph {
     pub total_allocation_count: OnceLock<u64>,
     pub corrected_self_time: OnceLock<u64>,
     pub corrected_total_time: OnceLock<u64>,
+    pub bottom_up: OnceLock<Vec<Arc<SpanBottomUp>>>,
 }
 
 pub struct SpanBottomUp {
     // These values won't change after creation:
     pub self_spans: Vec<SpanIndex>,
-    pub children: HashMap<String, SpanBottomUp>,
+    pub children: Vec<Arc<SpanBottomUp>>,
     pub example_span: SpanIndex,
 
     // These values are computed when accessed:
@@ -103,10 +104,14 @@ pub struct SpanBottomUp {
 }
 
 impl SpanBottomUp {
-    pub fn new(example_span: SpanIndex) -> Self {
+    pub fn new(
+        self_spans: Vec<SpanIndex>,
+        example_span: SpanIndex,
+        children: Vec<Arc<SpanBottomUp>>,
+    ) -> Self {
         Self {
-            self_spans: vec![],
-            children: HashMap::new(),
+            self_spans,
+            children,
             example_span,
             max_depth: OnceLock::new(),
             events: OnceLock::new(),

crates/turbopack-trace-server/src/span_bottom_up_ref.rs

@@ -1,6 +1,7 @@
 use std::{
     collections::VecDeque,
     fmt::{Debug, Formatter},
+    sync::Arc,
 };
 
 use indexmap::IndexMap;
@@ -13,7 +14,7 @@ use crate::{
 };
 
 pub struct SpanBottomUpRef<'a> {
-    pub(crate) bottom_up: &'a SpanBottomUp,
+    pub(crate) bottom_up: Arc<SpanBottomUp>,
     pub(crate) store: &'a Store,
 }
 
@@ -36,7 +37,7 @@ impl<'a> SpanBottomUpRef<'a> {
         }
     }
 
-    pub fn spans(&self) -> impl Iterator<Item = SpanRef<'a>> + 'a {
+    pub fn spans(&self) -> impl Iterator<Item = SpanRef<'a>> + '_ {
         let store = self.store;
         self.bottom_up.self_spans.iter().map(move |span| SpanRef {
             span: &store.spans[span.get()],
@@ -60,18 +61,18 @@ impl<'a> SpanBottomUpRef<'a> {
         }
     }
 
-    pub fn children(&self) -> impl Iterator<Item = SpanBottomUpRef<'a>> + 'a {
+    pub fn children(&self) -> impl Iterator<Item = SpanBottomUpRef<'a>> + '_ {
         self.bottom_up
             .children
-            .values()
+            .iter()
             .map(|bottom_up| SpanBottomUpRef {
-                bottom_up,
+                bottom_up: bottom_up.clone(),
                 store: self.store,
             })
     }
 
     #[allow(dead_code)]
-    pub fn graph(&self) -> impl Iterator<Item = SpanGraphEventRef<'a>> + 'a {
+    pub fn graph(&self) -> impl Iterator<Item = SpanGraphEventRef<'a>> + '_ {
         self.bottom_up
             .events
             .get_or_init(|| {

crates/turbopack-trace-server/src/span_graph_ref.rs

@@ -7,6 +7,7 @@ use std::{
 use indexmap::IndexMap;
 
 use crate::{
+    bottom_up::build_bottom_up_graph,
     span::{SpanGraph, SpanGraphEvent, SpanIndex},
     span_bottom_up_ref::SpanBottomUpRef,
     span_ref::SpanRef,
@@ -116,7 +117,14 @@ impl<'a> SpanGraphRef<'a> {
     }
 
     pub fn bottom_up(&self) -> impl Iterator<Item = SpanBottomUpRef<'a>> + '_ {
-        self.root_spans().flat_map(move |span| span.bottom_up())
+        self.graph
+            .bottom_up
+            .get_or_init(|| build_bottom_up_graph(self.root_spans()))
+            .iter()
+            .map(move |bottom_up| SpanBottomUpRef {
+                bottom_up: bottom_up.clone(),
+                store: self.store,
+            })
     }
 
     pub fn max_depth(&self) -> u32 {
@@ -269,6 +277,7 @@ pub fn event_map_to_list(
                 total_allocation_count: OnceLock::new(),
                 corrected_self_time: OnceLock::new(),
                 corrected_total_time: OnceLock::new(),
+                bottom_up: OnceLock::new(),
             };
             SpanGraphEvent::Child {
                 child: Arc::new(graph),

crates/turbopack-trace-server/src/span_ref.rs

@@ -283,14 +283,14 @@ impl<'a> SpanRef<'a> {
         self.span
             .bottom_up
             .get_or_init(|| {
-                build_bottom_up_graph(SpanRef {
+                build_bottom_up_graph([SpanRef {
                     span: self.span,
                     store: self.store,
-                })
+                }])
             })
             .iter()
             .map(|bottom_up| SpanBottomUpRef {
-                bottom_up,
+                bottom_up: bottom_up.clone(),
                 store: self.store,
             })
     }

crates/turbopack-trace-server/src/viewer.rs

@@ -369,7 +369,8 @@ impl Viewer {
                                 );
                             }
                         } else {
-                            let bottom_up = bottom_up.flat_map(|bottom_up| bottom_up.spans());
+                            let bottom_up = bottom_up
+                                .flat_map(|bottom_up| bottom_up.spans().collect::<Vec<_>>());
                             let bottom_up = if selected_view_mode.sort_children() {
                                 Either::Left(bottom_up.sorted_by_cached_key(|child| {
                                     Reverse(value_mode.value_from_bottom_up_span(child))
@@ -476,7 +477,8 @@ impl Viewer {
                                 );
                             }
                         } else {
-                            let bottom_up = bottom_up.flat_map(|bottom_up| bottom_up.spans());
+                            let bottom_up = bottom_up
+                                .flat_map(|bottom_up| bottom_up.spans().collect::<Vec<_>>());
                             let bottom_up = if selected_view_mode.sort_children() {
                                 Either::Left(bottom_up.sorted_by_cached_key(|child| {
                                     Reverse(value_mode.value_from_bottom_up_span(child))