Automatically expand and scroll the blueprint tree when focusing on an item (#5482)

### What

Double-clicking on an item, e.g. in a space views, focuses on that item.
When that happens, the blueprint tree UI will uncollapse and scroll as
needed to bring said item to view. When that item is an entity (e.g.
double-clicked in the streams view), then all occurrences of that entity
are uncollapsed, and the view scrolls to the first one.



https://github.com/rerun-io/rerun/assets/49431240/93b54fb6-ad52-4880-babd-d1001732c4db


### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested the web demo (if applicable):
* Using newly built examples:
[app.rerun.io](https://app.rerun.io/pr/5482/index.html)
* Using examples from latest `main` build:
[app.rerun.io](https://app.rerun.io/pr/5482/index.html?manifest_url=https://app.rerun.io/version/main/examples_manifest.json)
* Using full set of examples from `nightly` build:
[app.rerun.io](https://app.rerun.io/pr/5482/index.html?manifest_url=https://app.rerun.io/version/nightly/examples_manifest.json)
* [x] The PR title and labels are set such as to maximize their
usefulness for the next release's CHANGELOG
* [x] If applicable, add a new check to the [release
checklist](https://github.com/rerun-io/rerun/blob/main/tests/python/release_checklist)!

- [PR Build Summary](https://build.rerun.io/pr/5482)
- [Docs
preview](https://rerun.io/preview/ce5c582ffbe0844840846373a467a6907fb7b895/docs)
<!--DOCS-PREVIEW-->
- [Examples
preview](https://rerun.io/preview/ce5c582ffbe0844840846373a467a6907fb7b895/examples)
<!--EXAMPLES-PREVIEW-->
- [Recent benchmark results](https://build.rerun.io/graphs/crates.html)
- [Wasm size tracking](https://build.rerun.io/graphs/sizes.html)

Author: abey79

Cargo.lock

@@ -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/Cargo.toml

@@ -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/context_menu/actions/collapse_expand_all.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.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;
@@ -442,26 +443,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.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);
     }
 

crates/re_viewport/src/viewport_blueprint_ui.rs

@@ -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)