Automatically expand and scroll the blueprint tree when focusing on an item

crates/re_viewport/Cargo.toml

@@ -49,3 +49,4 @@ nohash-hasher.workspace = true
 once_cell.workspace = true
 rayon.workspace = true
 rmp-serde.workspace = true
+smallvec.workspace = true

crates/re_viewport/src/context_menu/actions/collapse_expand_all.rs

@@ -46,7 +46,7 @@ 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 {
+            .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()),

crates/re_viewport/src/viewport.rs

@@ -48,6 +48,13 @@ pub struct ViewportState {
     ///
     /// See [`ViewportState::is_candidate_drop_parent_container`] for details.
     candidate_drop_parent_container_id: Option<ContainerId>,
+
+    /// The item that should be focused on in the blueprint tree.
+    ///
+    /// Set at each frame by [`Viewport::tree_ui`]. This is similar to
+    /// [`ViewerContext::focused_item`] but account for how specifically the blueprint tree should
+    /// handle the focused item.
+    pub(crate) blueprint_tree_scroll_to_item: Option<Item>,
 }
 
 static DEFAULT_PROPS: Lazy<EntityPropertyMap> = Lazy::<EntityPropertyMap>::new(Default::default);

crates/re_viewport/src/viewport_blueprint.rs

@@ -3,8 +3,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 
 use ahash::HashMap;
 use egui_tiles::{SimplificationOptions, TileId};
-
 use nohash_hasher::IntSet;
+use smallvec::SmallVec;
+
 use re_data_store::LatestAtQuery;
 use re_entity_db::EntityPath;
 use re_log_types::hash::Hash64;
@@ -439,26 +440,50 @@ impl ViewportBlueprint {
             )
     }
 
+    /// Walk the entire [`Contents`] tree, starting from the root container.
+    ///
+    /// See [`Self::visit_contents_in_container`] for details.
+    pub fn visit_contents(&self, visitor: &mut impl FnMut(&Contents, &SmallVec<[ContainerId; 4]>)) {
+        if let Some(root_container) = self.root_container {
+            self.visit_contents_in_container(&root_container, visitor);
+        }
+    }
+
     /// Walk the subtree defined by the provided container id and call `visitor` for each
     /// [`Contents`].
     ///
-    /// Note: `visitor` is first called for the container passed in argument.
+    /// Note:
+    /// - `visitor` is first called for the container passed in argument
+    /// - `visitor`'s second argument contains the hierarchy leading to the visited contents, from
+    ///   (and including) the container passed in argument
     pub fn visit_contents_in_container(
         &self,
         container_id: &ContainerId,
-        visitor: &mut impl FnMut(&Contents),
+        visitor: &mut impl FnMut(&Contents, &SmallVec<[ContainerId; 4]>),
+    ) {
+        let mut hierarchy = SmallVec::new();
+        self.visit_contents_in_container_impl(container_id, &mut hierarchy, visitor);
+    }
+
+    fn visit_contents_in_container_impl(
+        &self,
+        container_id: &ContainerId,
+        hierarchy: &mut SmallVec<[ContainerId; 4]>,
+        visitor: &mut impl FnMut(&Contents, &SmallVec<[ContainerId; 4]>),
     ) {
-        visitor(&Contents::Container(*container_id));
+        visitor(&Contents::Container(*container_id), hierarchy);
         if let Some(container) = self.container(container_id) {
+            hierarchy.push(*container_id);
             for contents in &container.contents {
-                visitor(contents);
+                visitor(contents, hierarchy);
                 match contents {
                     Contents::Container(container_id) => {
-                        self.visit_contents_in_container(container_id, visitor);
+                        self.visit_contents_in_container_impl(container_id, hierarchy, visitor);
                     }
                     Contents::SpaceView(_) => {}
                 }
             }
+            hierarchy.pop();
         }
     }
 

crates/re_viewport/src/viewport_blueprint_ui.rs

@@ -1,6 +1,7 @@
 use egui::{Response, Ui};
 use itertools::Itertools;
 use re_data_ui::item_ui::guess_instance_path_icon;
+use smallvec::SmallVec;
 
 use re_entity_db::InstancePath;
 use re_log_types::EntityPath;
@@ -54,14 +55,16 @@ impl<'a> DataResultNodeOrPath<'a> {
 
 impl Viewport<'_, '_> {
     /// Show the blueprint panel tree view.
-    pub fn tree_ui(&self, ctx: &ViewerContext<'_>, ui: &mut egui::Ui) {
+    pub fn tree_ui(&mut self, ctx: &ViewerContext<'_>, ui: &mut egui::Ui) {
         re_tracing::profile_function!();
 
         egui::ScrollArea::both()
             .id_source("blueprint_tree_scroll_area")
             .auto_shrink([true, false])
             .show(ui, |ui| {
                 ctx.re_ui.panel_content(ui, |_, ui| {
+                    self.state.blueprint_tree_scroll_to_item = self.handle_focused_item(ctx, ui);
+
                     self.root_container_tree_ui(ctx, ui);
 
                     let empty_space_response =
@@ -81,6 +84,131 @@ impl Viewport<'_, '_> {
             });
     }
 
+    /// Expend all required items and compute which item we should scroll to.
+    fn handle_focused_item(&self, ctx: &ViewerContext<'_>, ui: &egui::Ui) -> Option<Item> {
+        let focused_item = ctx.focused_item.as_ref()?;
+        match focused_item {
+            Item::Container(container_id) => {
+                self.expand_all_contents_until(ui.ctx(), &Contents::Container(*container_id));
+                Some(focused_item.clone())
+            }
+            Item::SpaceView(space_view_id) => {
+                self.expand_all_contents_until(ui.ctx(), &Contents::SpaceView(*space_view_id));
+                ctx.focused_item.clone()
+            }
+            Item::DataResult(space_view_id, instance_path) => {
+                self.expand_all_contents_until(ui.ctx(), &Contents::SpaceView(*space_view_id));
+                self.expand_all_data_results_until(
+                    ctx,
+                    ui.ctx(),
+                    space_view_id,
+                    &instance_path.entity_path,
+                );
+
+                ctx.focused_item.clone()
+            }
+            Item::InstancePath(instance_path) => {
+                let space_view_ids =
+                    self.list_space_views_with_entity(ctx, &instance_path.entity_path);
+
+                // focus on the first matching data result
+                let res = space_view_ids
+                    .first()
+                    .map(|id| Item::DataResult(*id, instance_path.clone()));
+
+                for space_view_id in space_view_ids {
+                    self.expand_all_contents_until(ui.ctx(), &Contents::SpaceView(space_view_id));
+                    self.expand_all_data_results_until(
+                        ctx,
+                        ui.ctx(),
+                        &space_view_id,
+                        &instance_path.entity_path,
+                    );
+                }
+
+                res
+            }
+
+            Item::StoreId(_) | Item::ComponentPath(_) => None,
+        }
+    }
+
+    /// Expand all containers until reaching the provided content.
+    fn expand_all_contents_until(&self, egui_ctx: &egui::Context, focused_contents: &Contents) {
+        //TODO(ab): this could look nicer if `Contents` was declared in re_view_context :)
+        let expend_contents = |contents: &Contents| match contents {
+            Contents::Container(container_id) => CollapseScope::BlueprintTree
+                .container(*container_id)
+                .set_open(egui_ctx, true),
+            Contents::SpaceView(space_view_id) => CollapseScope::BlueprintTree
+                .space_view(*space_view_id)
+                .set_open(egui_ctx, true),
+        };
+
+        self.blueprint.visit_contents(&mut |contents, hierarchy| {
+            if contents == focused_contents {
+                expend_contents(contents);
+                for parent in hierarchy {
+                    expend_contents(&Contents::Container(*parent));
+                }
+            }
+        });
+    }
+
+    /// List all space views that have the provided entity as data result.
+    #[inline]
+    fn list_space_views_with_entity(
+        &self,
+        ctx: &ViewerContext<'_>,
+        entity_path: &EntityPath,
+    ) -> SmallVec<[SpaceViewId; 4]> {
+        let mut space_view_ids = SmallVec::new();
+        self.blueprint.visit_contents(&mut |contents, _| {
+            if let Contents::SpaceView(space_view_id) = contents {
+                let result_tree = &ctx.lookup_query_result(*space_view_id).tree;
+                if result_tree.lookup_node_by_path(entity_path).is_some() {
+                    space_view_ids.push(*space_view_id);
+                }
+            }
+        });
+        space_view_ids
+    }
+
+    /// Expand data results of the provided space view all the way to the provided entity.
+    #[allow(clippy::unused_self)]
+    fn expand_all_data_results_until(
+        &self,
+        ctx: &ViewerContext<'_>,
+        egui_ctx: &egui::Context,
+        space_view_id: &SpaceViewId,
+        entity_path: &EntityPath,
+    ) {
+        let result_tree = &ctx.lookup_query_result(*space_view_id).tree;
+        if result_tree.lookup_node_by_path(entity_path).is_some() {
+            if let Some(root_node) = result_tree.root_node() {
+                EntityPath::incremental_walk(Some(&root_node.data_result.entity_path), entity_path)
+                    .chain(std::iter::once(root_node.data_result.entity_path.clone()))
+                    .for_each(|entity_path| {
+                        CollapseScope::BlueprintTree
+                            .data_result(*space_view_id, entity_path)
+                            .set_open(egui_ctx, true);
+                    });
+            }
+        }
+    }
+
+    /// Check if the provided item should be scrolled to.
+    fn scroll_to_me_if_needed(&self, ui: &egui::Ui, item: &Item, response: &egui::Response) {
+        if Some(item) == self.state.blueprint_tree_scroll_to_item.as_ref() {
+            // Scroll only if the entity isn't already visible. This is important because that's what
+            // happens when double-clicking an entity _in the blueprint tree_. In such case, it would be
+            // annoying to induce a scroll motion.
+            if !ui.clip_rect().contains_rect(response.rect) {
+                response.scroll_to_me(Some(egui::Align::Center));
+            }
+        }
+    }
+
     /// If a group or spaceview has a total of this number of elements, show its subtree by default?
     fn default_open_for_data_result(group: &DataResultNode) -> bool {
         let num_children = group.children.len();
@@ -145,6 +273,7 @@ impl Viewport<'_, '_> {
             &item_response,
             SelectionUpdateBehavior::UseSelection,
         );
+        self.scroll_to_me_if_needed(ui, &item, &item_response);
         ctx.select_hovered_on_click(&item_response, item);
 
         self.handle_root_container_drag_and_drop_interaction(
@@ -218,6 +347,7 @@ impl Viewport<'_, '_> {
             &response,
             SelectionUpdateBehavior::UseSelection,
         );
+        self.scroll_to_me_if_needed(ui, &item, &response);
         ctx.select_hovered_on_click(&response, item);
 
         self.blueprint
@@ -304,15 +434,15 @@ impl Viewport<'_, '_> {
                     );
 
                     // Show 'projections' if there's any items that weren't part of the tree under origin but are directly included.
-                    // The later is important since `+ image/camera/**` necessarily has `image` and `image/camera` in the data result tree.
+                    // The latter is important since `+ image/camera/**` necessarily has `image` and `image/camera` in the data result tree.
                     let mut projections = Vec::new();
                     result_tree.visit(&mut |node| {
                         if node
                             .data_result
                             .entity_path
                             .starts_with(&space_view.space_origin)
                         {
-                            false // If its under the origin, we're not interested, stop recursing.
+                            false // If it's under the origin, we're not interested, stop recursing.
                         } else if node.data_result.tree_prefix_only {
                             true // Keep recursing until we find a projection.
                         } else {
@@ -351,6 +481,7 @@ impl Viewport<'_, '_> {
             &response,
             SelectionUpdateBehavior::UseSelection,
         );
+        self.scroll_to_me_if_needed(ui, &item, &response);
         ctx.select_hovered_on_click(&response, item);
 
         let content = Contents::SpaceView(*space_view_id);
@@ -474,8 +605,7 @@ impl Viewport<'_, '_> {
             let response = list_item
                 .show_collapsing(
                     ui,
-                    CollapseScope::BlueprintTree
-                        .data_result(space_view.id, node.data_result.entity_path.clone()),
+                    CollapseScope::BlueprintTree.data_result(space_view.id, entity_path.clone()),
                     default_open,
                     |_, ui| {
                         for child in node.children.iter().sorted_by_key(|c| {
@@ -526,6 +656,7 @@ impl Viewport<'_, '_> {
             &response,
             SelectionUpdateBehavior::UseSelection,
         );
+        self.scroll_to_me_if_needed(ui, &item, &response);
         ctx.select_hovered_on_click(&response, item);
     }
 

tests/python/release_checklist/check_focus.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import os
+from argparse import Namespace
+from uuid import uuid4
+
+import rerun as rr
+
+README = """
+# Focus checks
+
+## Preparation
+
+TODO(ab): automate this with blueprints
+TODO(ab): add lots of stuff via blueprint to make the tree more crowded and check scrolling
+
+- Reset the blueprint
+- Clone the 3D space view such as to have 2 of them.
+
+## Checks
+
+- Collapse all in the blueprint tree.
+- Double-click on the box in the first space view, check corresponding space view expands.
+- Collapse all in the blueprint tree.
+- Double-click on the leaf "boxes3d" entity in the streams view, check both space views expand.
+"""
+
+
+def log_readme() -> None:
+    rr.log("readme", rr.TextDocument(README, media_type=rr.MediaType.MARKDOWN), timeless=True)
+
+
+def log_some_space_views() -> None:
+    rr.set_time_sequence("frame_nr", 0)
+
+    rr.log(
+        "/objects/boxes/boxes3d",
+        rr.Boxes3D(centers=[[0, 0, 0], [1, 1.5, 1.15], [3, 2, 1]], half_sizes=[0.5, 1, 0.5] * 3),
+    )
+
+
+def run(args: Namespace) -> None:
+    # TODO(cmc): I have no idea why this works without specifying a `recording_id`, but
+    # I'm not gonna rely on it anyway.
+    rr.script_setup(args, f"{os.path.basename(__file__)}", recording_id=uuid4())
+
+    log_readme()
+    log_some_space_views()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Interactive release checklist")
+    rr.script_add_args(parser)
+    args = parser.parse_args()
+    run(args)