Bottom up tracing and some tracing improvements (vercel/turborepo#6923)

### Description

Add bottom up view to tracing
Tracing refactorings

### Testing Instructions

<!--
  Give a quick description of steps to test your changes.
-->


Closes PACK-2198

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

@@ -0,0 +1,78 @@
+use std::{collections::HashMap, sync::Arc};
+
+use either::Either;
+
+use crate::{
+    span::{SpanBottomUp, SpanIndex},
+    span_ref::SpanRef,
+};
+
+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![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() {
+            current_iterators.push(iter);
+
+            let name = child.group_name();
+            let (_, mut bottom_up) = roots
+                .raw_entry_mut()
+                .from_key(name)
+                .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() {
+                if prev == Some(name) {
+                    continue;
+                }
+                let (_, child_bottom_up) = bottom_up
+                    .children
+                    .raw_entry_mut()
+                    .from_key(name)
+                    .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(Either::Right(child.children()));
+        } else {
+            current_path.pop();
+        }
+    }
+    roots.into_values().map(|b| Arc::new(b.build())).collect()
+}

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

@@ -5,9 +5,13 @@ use std::{collections::HashSet, sync::Arc};
 
 use self::{reader::TraceReader, server::serve, store_container::StoreContainer};
 
+mod bottom_up;
 mod reader;
 mod server;
 mod span;
+mod span_bottom_up_ref;
+mod span_graph_ref;
+mod span_ref;
 mod store;
 mod store_container;
 mod u64_empty_string;

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

@@ -20,7 +20,7 @@ use crate::{
     store::SpanId,
     store_container::StoreContainer,
     u64_string,
-    viewer::{ViewLineUpdate, ViewMode, Viewer},
+    viewer::{Update, ViewLineUpdate, ViewMode, Viewer},
 };
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -33,6 +33,7 @@ pub enum ServerToClientMessage {
     },
     ViewLinesCount {
         count: usize,
+        max: u64,
     },
     #[serde(rename_all = "camelCase")]
     QueryResult {
@@ -44,6 +45,8 @@ pub enum ServerToClientMessage {
         duration: u64,
         allocations: u64,
         deallocations: u64,
+        allocation_count: u64,
+        persistent_allocations: u64,
         args: Vec<(String, String)>,
         path: Vec<String>,
     },
@@ -161,14 +164,17 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                         return Ok(());
                     }
                     state.last_update_generation = store.generation();
-                    let updates = state.viewer.compute_update(&store, &state.view_rect);
+                    let Update {
+                        lines: updates,
+                        max,
+                    } = state.viewer.compute_update(&store, &state.view_rect);
                     let count = updates.len();
                     for update in updates {
                         let message = ServerToClientMessage::ViewLine { update };
                         let message = serde_json::to_string(&message).unwrap();
                         state.writer.send_message(&OwnedMessage::Text(message))?;
                     }
-                    let message = ServerToClientMessage::ViewLinesCount { count };
+                    let message = ServerToClientMessage::ViewLinesCount { count, max };
                     let message = serde_json::to_string(&message).unwrap();
                     state.writer.send_message(&OwnedMessage::Text(message))?;
                     ready_for_update.store(false, Ordering::SeqCst);
@@ -212,39 +218,36 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                     )?;
                                 }
                                 ClientToServerMessage::ViewMode { id, mode, inherit } => {
-                                    match mode.as_str() {
+                                    let (mode, sorted) =
+                                        if let Some(mode) = mode.strip_suffix("-sorted") {
+                                            (mode, true)
+                                        } else {
+                                            (mode.as_str(), false)
+                                        };
+                                    match mode {
                                         "raw-spans" => {
                                             state.viewer.set_view_mode(
                                                 id,
-                                                Some((
-                                                    ViewMode::RawSpans { sorted: false },
-                                                    inherit,
-                                                )),
+                                                Some((ViewMode::RawSpans { sorted }, inherit)),
                                             );
                                         }
-                                        "raw-spans-sorted" => {
+                                        "aggregated" => {
                                             state.viewer.set_view_mode(
                                                 id,
-                                                Some((
-                                                    ViewMode::RawSpans { sorted: true },
-                                                    inherit,
-                                                )),
+                                                Some((ViewMode::Aggregated { sorted }, inherit)),
                                             );
                                         }
-                                        "aggregated" => {
+                                        "bottom-up" => {
                                             state.viewer.set_view_mode(
                                                 id,
-                                                Some((
-                                                    ViewMode::Aggregated { sorted: false },
-                                                    inherit,
-                                                )),
+                                                Some((ViewMode::BottomUp { sorted }, inherit)),
                                             );
                                         }
-                                        "aggregated-sorted" => {
+                                        "aggregated-bottom-up" => {
                                             state.viewer.set_view_mode(
                                                 id,
                                                 Some((
-                                                    ViewMode::Aggregated { sorted: true },
+                                                    ViewMode::AggregatedBottomUp { sorted },
                                                     inherit,
                                                 )),
                                             );
@@ -278,6 +281,9 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                         let duration = span.corrected_total_time();
                                         let allocations = span.total_allocations();
                                         let deallocations = span.total_deallocations();
+                                        let allocation_count = span.total_allocation_count();
+                                        let persistent_allocations =
+                                            span.total_persistent_allocations();
                                         let args = span
                                             .args()
                                             .map(|(k, v)| (k.to_string(), v.to_string()))
@@ -297,6 +303,8 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                             duration,
                                             allocations,
                                             deallocations,
+                                            allocation_count,
+                                            persistent_allocations,
                                             args,
                                             path,
                                         }
@@ -309,6 +317,8 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                                             duration: 0,
                                             allocations: 0,
                                             deallocations: 0,
+                                            allocation_count: 0,
+                                            persistent_allocations: 0,
                                             args: Vec::new(),
                                             path: Vec::new(),
                                         }

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

@@ -41,6 +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<Arc<SpanBottomUp>>>,
     pub search_index: OnceLock<HashMap<String, Vec<SpanIndex>>>,
 }
 
@@ -82,4 +83,44 @@ 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: Vec<Arc<SpanBottomUp>>,
+    pub example_span: SpanIndex,
+
+    // These values are computed when accessed:
+    pub max_depth: OnceLock<u32>,
+    pub events: OnceLock<Vec<SpanGraphEvent>>,
+    pub self_time: OnceLock<u64>,
+    pub corrected_self_time: OnceLock<u64>,
+    pub self_allocations: OnceLock<u64>,
+    pub self_deallocations: OnceLock<u64>,
+    pub self_persistent_allocations: OnceLock<u64>,
+    pub self_allocation_count: OnceLock<u64>,
+}
+
+impl SpanBottomUp {
+    pub fn new(
+        self_spans: Vec<SpanIndex>,
+        example_span: SpanIndex,
+        children: Vec<Arc<SpanBottomUp>>,
+    ) -> Self {
+        Self {
+            self_spans,
+            children,
+            example_span,
+            max_depth: OnceLock::new(),
+            events: OnceLock::new(),
+            self_time: OnceLock::new(),
+            corrected_self_time: OnceLock::new(),
+            self_allocations: OnceLock::new(),
+            self_deallocations: OnceLock::new(),
+            self_persistent_allocations: OnceLock::new(),
+            self_allocation_count: OnceLock::new(),
+        }
+    }
 }