Add keyboard navigation to the blueprint and streams tree (#9960)

### Related

* Part of #3055

### What

Now in both the blueprint tree and streams tree:
- left/right collapses/uncollapses the selected item
- up/down selects the previous/next item

**Notes**:
- This only works with single selection.
- This also fix a bug where shift-select could be broken by "collapsing
all" on the root container.
- Previously, the left/right arrows were used to move the time cursor.
This now requires cmd/alt-left/right. Also, the go to start/end of
timeline is now alt-left/right.

Author: abey79

crates/viewer/re_blueprint_tree/src/blueprint_tree.rs

@@ -1,3 +1,5 @@
+use std::ops::ControlFlow;
+
 use egui::{Response, Ui};
 use smallvec::SmallVec;
 
@@ -55,6 +57,12 @@ pub struct BlueprintTree {
     /// This is the item we used as a starting point for range selection. It is set and remembered
     /// everytime the user clicks on an item _without_ holding shift.
     range_selection_anchor_item: Option<Item>,
+
+    /// Used when the selection is modified using key navigation.
+    ///
+    /// IMPORTANT: Always make sure that the item will be drawn this or next frame when setting this
+    /// to `Some`, so that this flag is immediately consumed.
+    scroll_to_me_item: Option<Item>,
 }
 
 impl BlueprintTree {
@@ -186,6 +194,14 @@ impl BlueprintTree {
     ) {
         let item = Item::Container(container_data.id);
 
+        // It's possible that the root becomes technically collapsed (e.g. context menu or arrow
+        // navigation), even though we don't allow that in the ui. We really don't want that,
+        // though, because it breaks the collapse-based tree data visiting. To avoid that, we always
+        // force uncollapse this item.
+        self.collapse_scope()
+            .container(container_data.id)
+            .set_open(ctx.egui_ctx(), true);
+
         let item_response = ui
             .list_item()
             .render_offscreen(false)
@@ -218,7 +234,7 @@ impl BlueprintTree {
             viewport_blueprint,
             blueprint_tree_data,
             ui,
-            item,
+            &item,
             &item_response,
         );
 
@@ -336,7 +352,7 @@ impl BlueprintTree {
             viewport_blueprint,
             blueprint_tree_data,
             ui,
-            item,
+            &item,
             &response,
         );
 
@@ -448,7 +464,7 @@ impl BlueprintTree {
             viewport_blueprint,
             blueprint_tree_data,
             ui,
-            item,
+            &item,
             &response,
         );
 
@@ -618,7 +634,7 @@ impl BlueprintTree {
             viewport_blueprint,
             blueprint_tree_data,
             ui,
-            item,
+            &item,
             &response,
         );
     }
@@ -632,24 +648,119 @@ impl BlueprintTree {
         viewport_blueprint: &ViewportBlueprint,
         blueprint_tree_data: &BlueprintTreeData,
         ui: &egui::Ui,
-        item: Item,
+        item: &Item,
         response: &Response,
     ) {
         context_menu_ui_for_item_with_context(
             ctx,
             viewport_blueprint,
-            &item,
+            item,
             // expand/collapse context menu actions need this information
             ItemContext::BlueprintTree {
                 filter_session_id: self.filter_state.session_id(),
             },
             response,
             SelectionUpdateBehavior::UseSelection,
         );
-        self.scroll_to_me_if_needed(ui, &item, response);
+        self.scroll_to_me_if_needed(ui, item, response);
         ctx.handle_select_hover_drag_interactions(response, item.clone(), true);
 
-        self.handle_range_selection(ctx, blueprint_tree_data, item, response);
+        self.handle_range_selection(ctx, blueprint_tree_data, item.clone(), response);
+
+        self.handle_key_navigation(ctx, blueprint_tree_data, item);
+    }
+
+    fn handle_key_navigation(
+        &mut self,
+        ctx: &ViewerContext<'_>,
+        blueprint_tree_data: &BlueprintTreeData,
+        item: &Item,
+    ) {
+        if ctx.selection_state().selected_items().single_item() != Some(item) {
+            return;
+        }
+
+        if ctx
+            .egui_ctx()
+            .input_mut(|i| i.consume_key(egui::Modifiers::NONE, egui::Key::ArrowRight))
+        {
+            if let Some(collapse_id) = self.collapse_scope().item(item.clone()) {
+                collapse_id.set_open(ctx.egui_ctx(), true);
+            }
+        }
+
+        if ctx
+            .egui_ctx()
+            .input_mut(|i| i.consume_key(egui::Modifiers::NONE, egui::Key::ArrowLeft))
+        {
+            if let Some(collapse_id) = self.collapse_scope().item(item.clone()) {
+                collapse_id.set_open(ctx.egui_ctx(), false);
+            }
+        }
+
+        if ctx
+            .egui_ctx()
+            .input_mut(|i| i.consume_key(egui::Modifiers::NONE, egui::Key::ArrowDown))
+        {
+            let mut found_current = false;
+
+            let result = blueprint_tree_data.visit(|tree_item| {
+                let is_item_collapsed = !tree_item.is_open(ctx.egui_ctx(), self.collapse_scope());
+
+                if &tree_item.item() == item {
+                    found_current = true;
+
+                    return if is_item_collapsed {
+                        VisitorControlFlow::SkipBranch
+                    } else {
+                        VisitorControlFlow::Continue
+                    };
+                }
+
+                if found_current {
+                    VisitorControlFlow::Break(Some(tree_item.item()))
+                } else if is_item_collapsed {
+                    VisitorControlFlow::SkipBranch
+                } else {
+                    VisitorControlFlow::Continue
+                }
+            });
+
+            if let ControlFlow::Break(Some(item)) = result {
+                ctx.selection_state().set_selection(item.clone());
+                self.scroll_to_me_item = Some(item.clone());
+                self.range_selection_anchor_item = Some(item);
+            }
+        }
+
+        if ctx
+            .egui_ctx()
+            .input_mut(|i| i.consume_key(egui::Modifiers::NONE, egui::Key::ArrowUp))
+        {
+            let mut last_item = None;
+
+            let result = blueprint_tree_data.visit(|tree_item| {
+                let is_item_collapsed = !tree_item.is_open(ctx.egui_ctx(), self.collapse_scope());
+
+                if &tree_item.item() == item {
+                    return VisitorControlFlow::Break(last_item.clone());
+                }
+
+                last_item = Some(tree_item.item());
+
+                if is_item_collapsed {
+                    VisitorControlFlow::SkipBranch
+                } else {
+                    VisitorControlFlow::Continue
+                }
+            });
+
+            if let ControlFlow::Break(Some(item)) = result {
+                ctx.selection_state().set_selection(item.clone());
+                self.scroll_to_me_item = Some(item.clone());
+                self.range_selection_anchor_item = Some(item);
+            }
+        }
     }
 
     /// Handle setting/extending the selection based on shift-clicking.
@@ -738,9 +849,7 @@ impl BlueprintTree {
                 return VisitorControlFlow::Break(());
             }
 
-            let is_expanded = blueprint_tree_item
-                .is_open(ctx.egui_ctx(), collapse_scope)
-                .unwrap_or(false);
+            let is_expanded = blueprint_tree_item.is_open(ctx.egui_ctx(), collapse_scope);
 
             if is_expanded {
                 VisitorControlFlow::Continue
@@ -757,7 +866,7 @@ impl BlueprintTree {
     }
 
     /// Check if the provided item should be scrolled to.
-    fn scroll_to_me_if_needed(&self, ui: &egui::Ui, item: &Item, response: &egui::Response) {
+    fn scroll_to_me_if_needed(&mut self, ui: &egui::Ui, item: &Item, response: &egui::Response) {
         if Some(item) == self.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
@@ -766,6 +875,13 @@ impl BlueprintTree {
                 response.scroll_to_me(Some(egui::Align::Center));
             }
         }
+
+        if Some(item) == self.scroll_to_me_item.as_ref() {
+            // This is triggered by keyboard navigation, so in this case we just want to scroll
+            // minimally for the item to be visible.
+            response.scroll_to_me(None);
+            self.scroll_to_me_item = None;
+        }
     }
 
     // ----------------------------------------------------------------------------

crates/viewer/re_blueprint_tree/src/data.rs

@@ -628,8 +628,8 @@ impl BlueprintTreeItem<'_> {
         }
     }
 
-    pub fn is_open(&self, ctx: &egui::Context, collapse_scope: CollapseScope) -> Option<bool> {
-        collapse_scope.item(self.item()).map(|collapse_id| {
+    pub fn is_open(&self, ctx: &egui::Context, collapse_scope: CollapseScope) -> bool {
+        collapse_scope.item(self.item()).is_some_and(|collapse_id| {
             collapse_id
                 .is_open(ctx)
                 .unwrap_or_else(|| self.default_open())

crates/viewer/re_time_panel/src/streams_tree_data.rs

@@ -7,7 +7,7 @@ use smallvec::SmallVec;
 use re_chunk_store::ChunkStore;
 use re_data_ui::sorted_component_list_for_ui;
 use re_entity_db::{EntityTree, InstancePath};
-use re_log_types::EntityPath;
+use re_log_types::{ComponentPath, EntityPath};
 use re_ui::filter_widget::{FilterMatcher, PathRanges};
 use re_viewer_context::{CollapseScope, Item, ViewerContext, VisitorControlFlow};
 
@@ -61,17 +61,17 @@ impl StreamsTreeData {
     /// Visit the entire tree.
     ///
     /// Note that we ALSO visit components, despite them not being part of the data structures. This
-    /// is because _currently_, we rarely need to visit, but when we do, we need to components, and
+    /// is because _currently_, we rarely need to visit, but when we do, we need components, and
     /// having them in the structure would be too expensive for the cases where it's unnecessary
     /// (e.g., when the tree is collapsed).
     ///
-    /// The provided closure is called once for each entity with `None` as component name argument.
+    /// The provided closure is called once for each entity with `None` as component argument.
     /// Then, consistent with the display order, its children entities are visited, and then its
     /// components are visited.
     pub fn visit<B>(
         &self,
         entity_db: &re_entity_db::EntityDb,
-        mut visitor: impl FnMut(&EntityData, Option<ComponentDescriptor>) -> VisitorControlFlow<B>,
+        mut visitor: impl FnMut(EntityOrComponentData<'_>) -> VisitorControlFlow<B>,
     ) -> ControlFlow<B> {
         let engine = entity_db.storage_engine();
         let store = engine.store();
@@ -214,16 +214,20 @@ impl EntityData {
     pub fn visit<B>(
         &self,
         store: &ChunkStore,
-        visitor: &mut impl FnMut(&Self, Option<ComponentDescriptor>) -> VisitorControlFlow<B>,
+        visitor: &mut impl FnMut(EntityOrComponentData<'_>) -> VisitorControlFlow<B>,
     ) -> ControlFlow<B> {
-        if visitor(self, None).visit_children()? {
+        if visitor(EntityOrComponentData::Entity(self)).visit_children()? {
             for child in &self.children {
                 child.visit(store, visitor)?;
             }
 
-            for component_descr in components_for_entity(store, &self.entity_path) {
+            for component_descriptor in components_for_entity(store, &self.entity_path) {
                 // these cannot have children
-                let _ = visitor(self, Some(component_descr)).visit_children()?;
+                let _ = visitor(EntityOrComponentData::Component {
+                    entity_data: self,
+                    component_descriptor,
+                })
+                .visit_children()?;
             }
         }
 
@@ -234,10 +238,10 @@ impl EntityData {
         Item::InstancePath(InstancePath::entity_all(self.entity_path.clone()))
     }
 
-    pub fn is_open(&self, ctx: &egui::Context, collapse_scope: CollapseScope) -> Option<bool> {
+    pub fn is_open(&self, ctx: &egui::Context, collapse_scope: CollapseScope) -> bool {
         collapse_scope
             .item(self.item())
-            .map(|collapse_id| collapse_id.is_open(ctx).unwrap_or(self.default_open))
+            .is_some_and(|collapse_id| collapse_id.is_open(ctx).unwrap_or(self.default_open))
     }
 }
 
@@ -252,3 +256,36 @@ pub fn components_for_entity(
         itertools::Either::Right(std::iter::empty())
     }
 }
+
+// ---
+
+#[derive(Debug)]
+pub enum EntityOrComponentData<'a> {
+    Entity(&'a EntityData),
+    Component {
+        entity_data: &'a EntityData,
+        component_descriptor: ComponentDescriptor,
+    },
+}
+
+impl EntityOrComponentData<'_> {
+    pub fn item(&self) -> Item {
+        match self {
+            Self::Entity(entity_data) => entity_data.item(),
+            Self::Component {
+                entity_data,
+                component_descriptor,
+            } => Item::ComponentPath(ComponentPath::new(
+                entity_data.entity_path.clone(),
+                component_descriptor.clone(),
+            )),
+        }
+    }
+
+    pub fn is_open(&self, ctx: &egui::Context, collapse_scope: CollapseScope) -> bool {
+        match self {
+            Self::Entity(entity_data) => entity_data.is_open(ctx, collapse_scope),
+            Self::Component { .. } => true,
+        }
+    }
+}