From 924141465c8de6c7b9925cc6705224aca0878e22 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 | 45 ++++++++++++++++++-
 crates/turbopack-trace-server/src/viewer.rs | 49 +++++++++++++++------
 2 files changed, 79 insertions(+), 15 deletions(-)

diff --git a/crates/turbopack-trace-server/src/server.rs b/crates/turbopack-trace-server/src/server.rs
index b579bdf58bbb24..849a2b1429eca9 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::{Shutdown, TcpStream},
+    num::NonZeroUsize,
     sync::{
         atomic::{AtomicBool, Ordering},
         Arc, Mutex,
@@ -17,10 +19,12 @@ use websocket::{
 };
 
 use crate::{
+    span::Span,
+    span_ref::SpanRef,
     store::SpanId,
     store_container::StoreContainer,
     u64_string,
-    viewer::{Update, ViewLineUpdate, ViewMode, Viewer},
+    viewer::{Update, ViewLineUpdate, ViewMode, ViewSpan, Viewer},
 };
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -86,6 +90,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 {
@@ -97,6 +114,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 {
@@ -141,6 +182,8 @@ pub fn serve(store: Arc<StoreContainer>) -> Result<()> {
                         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 8e3790fbdb7a51..342b1a0b4de24f 100644
--- a/crates/turbopack-trace-server/src/viewer.rs
+++ b/crates/turbopack-trace-server/src/viewer.rs
@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     server::ViewRect,
     span_bottom_up_ref::SpanBottomUpRef,
-    span_graph_ref::{SpanGraphEventRef, SpanGraphRef},
+    span_graph_ref::{self, SpanGraphEventRef, SpanGraphRef},
     span_ref::SpanRef,
     store::{SpanId, Store},
     u64_empty_string,
@@ -429,18 +429,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 {
+                if view_rect.should_filter_ref(&span, &mut highlighted_spans) {
                     children.last_mut().unwrap().item.filtered = true;
                 }
             }
@@ -455,7 +444,7 @@ impl Viewer {
             start,
             placeholder,
             view_mode,
-            filtered,
+            mut filtered,
         }) = queue.pop()
         {
             let line = get_line(&mut lines, line_index);
@@ -767,6 +756,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,
@@ -833,6 +852,7 @@ impl Viewer {
                         }
                         LineEntryType::SpanGraph(graph, filtered) => {
                             let (category, text) = graph.nice_name();
+
                             ViewSpan {
                                 id: graph.id().get() as u64,
                                 start: entry.start,
@@ -847,6 +867,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,