Filter entities in the UI (part 4): Add entity filtering in the blueprint tree

crates/viewer/re_context_menu/src/actions/collapse_expand_all.rs

@@ -1,13 +1,21 @@
+//! Context menu action to expand and collapse various trees in the UI.
+//!
+//! Note: the actual collapse/expand logic is in [`crate::collapse_expand`].
+
 use re_entity_db::InstancePath;
 use re_log_types::StoreKind;
-use re_viewer_context::{CollapseScope, ContainerId, Contents, Item, ItemContext, ViewId};
+use re_viewer_context::{CollapseScope, ContainerId, Item, ItemContext, ViewId};
 
+use crate::collapse_expand::{
+    collapse_expand_container, collapse_expand_data_result, collapse_expand_instance_path,
+    collapse_expand_view,
+};
 use crate::{ContextMenuAction, ContextMenuContext};
 
 /// Collapse or expand all items in the selection.
-// TODO(ab): the current implementation makes strong assumptions of which CollapseScope to use based
-// on the item type. This is brittle and will not scale if/when we add more trees to the UI. When
-// that happens, we will have to pass the scope to `context_menu_ui_for_item` and use it here.
+///
+/// Note: this makes _heavy_ use of [`ItemContext`] to determine the correct scope to
+/// collapse/expand.
 pub(crate) enum CollapseExpandAllAction {
     CollapseAll,
     ExpandAll,
@@ -50,29 +58,21 @@ impl ContextMenuAction for CollapseExpandAllAction {
     }
 
     fn process_container(&self, ctx: &ContextMenuContext<'_>, container_id: &ContainerId) {
-        ctx.viewport_blueprint
-            .visit_contents_in_container(container_id, &mut |contents, _| match contents {
-                Contents::Container(container_id) => CollapseScope::BlueprintTree
-                    .container(*container_id)
-                    .set_open(&ctx.egui_context, self.open()),
-                Contents::View(view_id) => self.process_view(ctx, view_id),
-            });
+        let scope = blueprint_collapse_scope(ctx, &Item::Container(*container_id));
+
+        collapse_expand_container(
+            ctx.viewer_context,
+            ctx.viewport_blueprint,
+            container_id,
+            scope,
+            self.open(),
+        );
     }
 
     fn process_view(&self, ctx: &ContextMenuContext<'_>, view_id: &ViewId) {
-        CollapseScope::BlueprintTree
-            .view(*view_id)
-            .set_open(&ctx.egui_context, self.open());
-
-        let query_result = ctx.viewer_context.lookup_query_result(*view_id);
-        let result_tree = &query_result.tree;
-        if let Some(root_node) = result_tree.root_node() {
-            self.process_data_result(
-                ctx,
-                view_id,
-                &InstancePath::entity_all(root_node.data_result.entity_path.clone()),
-            );
-        }
+        let scope = blueprint_collapse_scope(ctx, &Item::View(*view_id));
+
+        collapse_expand_view(ctx.viewer_context, view_id, scope, self.open());
     }
 
     fn process_data_result(
@@ -81,34 +81,23 @@ impl ContextMenuAction for CollapseExpandAllAction {
         view_id: &ViewId,
         instance_path: &InstancePath,
     ) {
-        //TODO(ab): here we should in principle walk the DataResult tree instead of the entity tree
-        // but the current API isn't super ergonomic.
-        let Some(subtree) = ctx
-            .viewer_context
-            .recording()
-            .tree()
-            .subtree(&instance_path.entity_path)
-        else {
-            return;
-        };
-
-        subtree.visit_children_recursively(|entity_path| {
-            CollapseScope::BlueprintTree
-                .data_result(*view_id, entity_path.clone())
-                .set_open(&ctx.egui_context, self.open());
-        });
+        let scope =
+            blueprint_collapse_scope(ctx, &Item::DataResult(*view_id, instance_path.clone()));
+
+        collapse_expand_data_result(
+            ctx.viewer_context,
+            view_id,
+            instance_path,
+            scope,
+            self.open(),
+        );
     }
 
     fn process_instance_path(&self, ctx: &ContextMenuContext<'_>, instance_path: &InstancePath) {
         let (db, scope) = match ctx
             .selection
             .context_for_item(&Item::InstancePath(instance_path.clone()))
         {
-            Some(&ItemContext::StreamsTree {
-                store_kind: StoreKind::Recording,
-                filter_session_id: None,
-            }) => (ctx.viewer_context.recording(), CollapseScope::StreamsTree),
-
             Some(&ItemContext::StreamsTree {
                 store_kind: StoreKind::Recording,
                 filter_session_id: Some(session_id),
@@ -133,21 +122,21 @@ impl ContextMenuAction for CollapseExpandAllAction {
                 CollapseScope::BlueprintStreamsTreeFiltered { session_id },
             ),
 
-            // default to recording if we don't have more specific information
-            Some(&ItemContext::TwoD { .. } | &ItemContext::ThreeD { .. }) | None => {
-                (ctx.viewer_context.recording(), CollapseScope::StreamsTree)
-            }
-        };
-
-        let Some(subtree) = db.tree().subtree(&instance_path.entity_path) else {
-            return;
+            // default to recording if the item context explicitly points to it or if we don't have
+            // any relevant context
+            Some(
+                &ItemContext::StreamsTree {
+                    store_kind: StoreKind::Recording,
+                    filter_session_id: None,
+                }
+                | &ItemContext::TwoD { .. }
+                | &ItemContext::ThreeD { .. }
+                | &ItemContext::BlueprintTree { .. },
+            )
+            | None => (ctx.viewer_context.recording(), CollapseScope::StreamsTree),
         };
 
-        subtree.visit_children_recursively(|entity_path| {
-            scope
-                .entity(entity_path.clone())
-                .set_open(&ctx.egui_context, self.open());
-        });
+        collapse_expand_instance_path(ctx.viewer_context, db, instance_path, scope, self.open());
     }
 }
 
@@ -159,3 +148,22 @@ impl CollapseExpandAllAction {
         }
     }
 }
+
+/// Determine the [`CollapseScope`] to use for items in the blueprint tree.
+fn blueprint_collapse_scope(ctx: &ContextMenuContext<'_>, item: &Item) -> CollapseScope {
+    match ctx.selection.context_for_item(item) {
+        Some(&ItemContext::BlueprintTree {
+            filter_session_id: Some(session_id),
+        }) => CollapseScope::BlueprintTreeFiltered { session_id },
+
+        None
+        | Some(
+            &ItemContext::BlueprintTree {
+                filter_session_id: None,
+            }
+            | &ItemContext::StreamsTree { .. }
+            | &ItemContext::TwoD { .. }
+            | &ItemContext::ThreeD { .. },
+        ) => CollapseScope::BlueprintTree,
+    }
+}

crates/viewer/re_context_menu/src/collapse_expand.rs

@@ -0,0 +1,88 @@
+//! Logic that implements the collapse/expand functionality.
+//!
+//! This is separated from the corresponding context menu action, so it may be reused directly, in
+//! particular in tests.
+
+use re_entity_db::{EntityDb, InstancePath};
+use re_viewer_context::{CollapseScope, ContainerId, Contents, ViewId, ViewerContext};
+use re_viewport_blueprint::ViewportBlueprint;
+
+pub fn collapse_expand_container(
+    ctx: &ViewerContext<'_>,
+    blueprint: &ViewportBlueprint,
+    container_id: &ContainerId,
+    scope: CollapseScope,
+    expand: bool,
+) {
+    blueprint.visit_contents_in_container(container_id, &mut |contents, _| {
+        match contents {
+            Contents::Container(container_id) => scope
+                .container(*container_id)
+                .set_open(ctx.egui_ctx, expand),
+
+            // IMPORTANT: don't call process_view() here, or the scope information would be lost
+            Contents::View(view_id) => collapse_expand_view(ctx, view_id, scope, expand),
+        }
+
+        true
+    });
+}
+
+pub fn collapse_expand_view(
+    ctx: &ViewerContext<'_>,
+    view_id: &ViewId,
+    scope: CollapseScope,
+    expand: bool,
+) {
+    scope.view(*view_id).set_open(ctx.egui_ctx, expand);
+
+    let query_result = ctx.lookup_query_result(*view_id);
+    let result_tree = &query_result.tree;
+    if let Some(root_node) = result_tree.root_node() {
+        collapse_expand_data_result(
+            ctx,
+            view_id,
+            &InstancePath::entity_all(root_node.data_result.entity_path.clone()),
+            scope,
+            expand,
+        );
+    }
+}
+
+pub fn collapse_expand_data_result(
+    ctx: &ViewerContext<'_>,
+    view_id: &ViewId,
+    instance_path: &InstancePath,
+    scope: CollapseScope,
+    expand: bool,
+) {
+    //TODO(ab): here we should in principle walk the DataResult tree instead of the entity tree
+    // but the current API isn't super ergonomic.
+    let Some(subtree) = ctx.recording().tree().subtree(&instance_path.entity_path) else {
+        return;
+    };
+
+    subtree.visit_children_recursively(|entity_path| {
+        scope
+            .data_result(*view_id, entity_path.clone())
+            .set_open(ctx.egui_ctx, expand);
+    });
+}
+
+pub fn collapse_expand_instance_path(
+    ctx: &ViewerContext<'_>,
+    db: &EntityDb,
+    instance_path: &InstancePath,
+    scope: CollapseScope,
+    expand: bool,
+) {
+    let Some(subtree) = db.tree().subtree(&instance_path.entity_path) else {
+        return;
+    };
+
+    subtree.visit_children_recursively(|entity_path| {
+        scope
+            .entity(entity_path.clone())
+            .set_open(ctx.egui_ctx, expand);
+    });
+}

crates/viewer/re_context_menu/src/lib.rs

@@ -9,6 +9,7 @@ use re_viewer_context::{
 use re_viewport_blueprint::{ContainerBlueprint, ViewportBlueprint};
 
 mod actions;
+pub mod collapse_expand;
 mod sub_menu;
 
 use actions::{

crates/viewer/re_time_panel/src/lib.rs

@@ -183,6 +183,11 @@ impl TimePanel {
         }
     }
 
+    /// Activates the search filter (for e.g. test purposes).
+    pub fn activate_filter(&mut self, query: &str) {
+        self.filter_state.activate(query);
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn show_panel(
         &mut self,
@@ -446,7 +451,7 @@ impl TimePanel {
 
             ui.add_space(-4.0); // hack to vertically center the text
 
-            let size = egui::vec2(self.prev_col_width, 28.0);
+            let size = egui::vec2(self.prev_col_width, 27.0);
             ui.allocate_ui_with_layout(size, egui::Layout::top_down(egui::Align::LEFT), |ui| {
                 ui.set_min_size(size);
                 ui.style_mut().wrap_mode = Some(egui::TextWrapMode::Extend);
@@ -463,8 +468,6 @@ impl TimePanel {
                         .strong(),
                     );
                 });
-
-                ui.add_space(2.0); // hack to vertically center the text
             })
             .response
             .on_hover_text(
@@ -749,7 +752,7 @@ impl TimePanel {
         let id = collapse_scope.entity(tree.path.clone()).into();
 
         let is_short_path = tree.path.len() <= 1 && !tree.is_leaf();
-        let default_open = self.filter_state.session_id().is_some() || is_short_path;
+        let default_open = self.filter_state.is_active() || is_short_path;
 
         let list_item::ShowCollapsingResponse {
             item_response: response,