From 7a20d4d8a38c3abe29122bd2114149840957b31b Mon Sep 17 00:00:00 2001
From: Alexander Lyon <arlyon@me.com>
Date: Fri, 8 Mar 2024 08:05:34 +0000
Subject: [PATCH] add ability to filter by value and occurences

---
 crates/turbopack-trace-server/src/server.rs | 42 +++++++++++++++++
 crates/turbopack-trace-server/src/viewer.rs | 51 +++++++++++++++------
 2 files changed, 78 insertions(+), 15 deletions(-)

diff --git a/crates/turbopack-trace-server/src/server.rs b/crates/turbopack-trace-server/src/server.rs
index 97fbf2e0cdcbf9..aed217a9f53920 100644
--- a/crates/turbopack-trace-server/src/server.rs
+++ b/crates/turbopack-trace-server/src/server.rs
@@ -1,5 +1,7 @@
 use std::{
+    collections::HashSet,
     net::{TcpListener, TcpStream},
+    num::NonZeroUsize,
     sync::{Arc, Mutex},
     thread::spawn,
 };
@@ -9,6 +11,7 @@ use serde::{Deserialize, Serialize};
 use tungstenite::{accept, Message};
 
 use crate::{
+    span_ref::SpanRef,
     store::SpanId,
     store_container::StoreContainer,
     u64_string,
@@ -80,6 +83,19 @@ pub struct SpanViewEvent {
     pub id: Option<SpanId>,
 }
 
+#[derive(Serialize, Deserialize, Debug)]
+pub struct Filter {
+    pub op: Op,
+    pub value: u64,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+#[serde(rename_all = "snake_case")]
+pub enum Op {
+    Gt,
+    Lt,
+}
+
 #[derive(Serialize, Deserialize, Debug)]
 #[serde(rename_all = "camelCase")]
 pub struct ViewRect {
@@ -91,6 +107,30 @@ pub struct ViewRect {
     pub query: String,
     pub view_mode: String,
     pub value_mode: String,
+    pub value_filter: Option<Filter>,
+    pub count_filter: Option<Filter>,
+}
+
+impl ViewRect {
+    pub fn should_filter_ref(
+        &self,
+        span: &SpanRef,
+        highlighted_spans: &mut HashSet<NonZeroUsize>,
+    ) -> bool {
+        let mut has_results = false;
+        for mut result in span.search(&self.query) {
+            has_results = true;
+            highlighted_spans.insert(result.id());
+            while let Some(parent) = result.parent() {
+                result = parent;
+                if !highlighted_spans.insert(result.id()) {
+                    break;
+                }
+            }
+        }
+
+        !has_results
+    }
 }
 
 struct ConnectionState {
@@ -130,6 +170,8 @@ fn handle_connection(
             query: String::new(),
             view_mode: "aggregated".to_string(),
             value_mode: "duration".to_string(),
+            count_filter: None,
+            value_filter: None,
         },
         last_update_generation: 0,
     }));
diff --git a/crates/turbopack-trace-server/src/viewer.rs b/crates/turbopack-trace-server/src/viewer.rs
index fc8f6ac181bc00..a163e473d8684e 100644
--- a/crates/turbopack-trace-server/src/viewer.rs
+++ b/crates/turbopack-trace-server/src/viewer.rs
@@ -47,6 +47,8 @@ impl ValueMode {
             ValueMode::PersistentAllocations => ValueMode::Allocations,
             ValueMode::AllocationCount => ValueMode::Allocations,
             ValueMode::Count => ValueMode::Count,
+            ValueMode::AllocationsPerTime => ValueMode::Allocations,
+            ValueMode::AllocationCountPerTime => ValueMode::Allocations,
         }
     }
 
@@ -471,20 +473,7 @@ impl Viewer {
                 false,
             ) && search_mode
             {
-                let mut has_results = false;
-                for mut result in span.search(&view_rect.query) {
-                    has_results = true;
-                    highlighted_spans.insert(result.id());
-                    while let Some(parent) = result.parent() {
-                        result = parent;
-                        if !highlighted_spans.insert(result.id()) {
-                            break;
-                        }
-                    }
-                }
-                if has_results {
-                    highlighted_spans.insert(span.id());
-                } else {
+                if view_rect.should_filter_ref(&span, &mut highlighted_spans) {
                     children.last_mut().unwrap().item.filtered = true;
                 }
             }
@@ -499,7 +488,7 @@ impl Viewer {
             start,
             placeholder,
             view_mode,
-            filtered,
+            mut filtered,
         }) = queue.pop()
         {
             let line = get_line(&mut lines, line_index);
@@ -823,6 +812,36 @@ impl Viewer {
                 // add children to queue
                 enqueue_children(children, &mut queue);
 
+                // check if we should filter based on width or count
+                if !filtered {
+                    let count = match &span {
+                        QueueItem::Span(_) => 1,
+                        QueueItem::SpanGraph(span_graph) => span_graph.count(),
+                        QueueItem::SpanBottomUp(bottom_up) => bottom_up.count(),
+                        QueueItem::SpanBottomUpSpan(_) => 1,
+                    };
+
+                    let emit = view_rect
+                        .count_filter
+                        .as_ref()
+                        .map(|filter| match filter.op {
+                            crate::server::Op::Gt => count > filter.value as usize,
+                            crate::server::Op::Lt => count < filter.value as usize,
+                        })
+                        .unwrap_or(true);
+
+                    let emit2 = view_rect
+                        .value_filter
+                        .as_ref()
+                        .map(|filter| match filter.op {
+                            crate::server::Op::Gt => width > filter.value,
+                            crate::server::Op::Lt => width < filter.value,
+                        })
+                        .unwrap_or(true);
+
+                    filtered = !emit || !emit2;
+                }
+
                 // add span to line
                 line.push(LineEntry {
                     start,
@@ -898,6 +917,7 @@ impl Viewer {
                         }
                         LineEntryType::SpanGraph(graph, filtered) => {
                             let (category, text) = graph.nice_name();
+
                             ViewSpan {
                                 id: graph.id().get() as u64,
                                 start: entry.start,
@@ -913,6 +933,7 @@ impl Viewer {
                         }
                         LineEntryType::SpanBottomUp(bottom_up, filtered) => {
                             let (category, text) = bottom_up.nice_name();
+
                             ViewSpan {
                                 id: bottom_up.id().get() as u64,
                                 start: entry.start,