Add the ability to provide local data to context menu actions

Cargo.lock

@@ -5782,6 +5782,7 @@ dependencies = [
  "re_viewer_context",
  "re_viewport_blueprint",
  "static_assertions",
+ "type-map",
 ]
 
 [[package]]

crates/viewer/re_context_menu/Cargo.toml

@@ -35,3 +35,4 @@ itertools.workspace = true
 nohash-hasher.workspace = true
 once_cell.workspace = true
 static_assertions.workspace = true
+type-map.workspace = true

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

@@ -49,17 +49,24 @@ impl ContextMenuAction for CollapseExpandAllAction {
     }
 
     fn process_container(&self, ctx: &ContextMenuContext<'_>, container_id: &ContainerId) {
+        let collapse_scope = ctx
+            .local_date()
+            .copied()
+            .unwrap_or(CollapseScope::BlueprintTree);
+
         ctx.viewport_blueprint
             .visit_contents_in_container(container_id, &mut |contents, _| match contents {
-                Contents::Container(container_id) => CollapseScope::BlueprintTree
+                Contents::Container(container_id) => collapse_scope
                     .container(*container_id)
                     .set_open(&ctx.egui_context, self.open()),
                 Contents::View(view_id) => self.process_view(ctx, view_id),
             });
     }
 
     fn process_view(&self, ctx: &ContextMenuContext<'_>, view_id: &ViewId) {
-        CollapseScope::BlueprintTree
+        ctx.local_date()
+            .copied()
+            .unwrap_or(CollapseScope::BlueprintTree)
             .view(*view_id)
             .set_open(&ctx.egui_context, self.open());
 
@@ -82,34 +89,34 @@ impl ContextMenuAction for CollapseExpandAllAction {
     ) {
         //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 {
+        let Some(subtree) = get_entity_tree(ctx, instance_path) else {
             return;
         };
 
+        let collapse_scope = ctx
+            .local_date()
+            .copied()
+            .unwrap_or(CollapseScope::BlueprintTree);
+
         subtree.visit_children_recursively(|entity_path| {
-            CollapseScope::BlueprintTree
+            collapse_scope
                 .data_result(*view_id, entity_path.clone())
                 .set_open(&ctx.egui_context, self.open());
         });
     }
 
     fn process_instance_path(&self, ctx: &ContextMenuContext<'_>, instance_path: &InstancePath) {
-        let Some(subtree) = ctx
-            .viewer_context
-            .recording()
-            .tree()
-            .subtree(&instance_path.entity_path)
-        else {
+        let Some(subtree) = get_entity_tree(ctx, instance_path) else {
             return;
         };
 
+        let collapse_scope = ctx
+            .local_date()
+            .copied()
+            .unwrap_or(CollapseScope::StreamsTree);
+
         subtree.visit_children_recursively(|entity_path| {
-            CollapseScope::StreamsTree
+            collapse_scope
                 .entity(entity_path.clone())
                 .set_open(&ctx.egui_context, self.open());
         });
@@ -124,3 +131,24 @@ impl CollapseExpandAllAction {
         }
     }
 }
+
+/// Get an [`re_entity_db::EntityTree`] for the given instance path.
+///
+/// This function guesses which store to search the entity in based on the [`CollapseScope`], which
+/// may be overridden as local data by the user code.
+fn get_entity_tree<'a>(
+    ctx: &'_ ContextMenuContext<'a>,
+    instance_path: &InstancePath,
+) -> Option<&'a re_entity_db::EntityTree> {
+    let collapse_scope = ctx
+        .local_date()
+        .copied()
+        .unwrap_or(CollapseScope::StreamsTree);
+
+    match collapse_scope {
+        CollapseScope::StreamsTree | CollapseScope::BlueprintTree => ctx.viewer_context.recording(),
+        CollapseScope::BlueprintStreamsTree => ctx.viewer_context.blueprint_db(),
+    }
+    .tree()
+    .subtree(&instance_path.entity_path)
+}

crates/viewer/re_context_menu/src/lib.rs

@@ -21,6 +21,10 @@ use actions::{
 };
 use sub_menu::SubMenu;
 
+pub mod exports {
+    pub use type_map;
+}
+
 /// Controls how [`context_menu_ui_for_item`] should handle the current selection state.
 #[derive(Debug, Clone, Copy)]
 pub enum SelectionUpdateBehavior {
@@ -41,6 +45,51 @@ pub fn context_menu_ui_for_item(
     item: &Item,
     item_response: &egui::Response,
     selection_update_behavior: SelectionUpdateBehavior,
+) {
+    context_menu_ui_for_item_with_local_data_map(
+        ctx,
+        viewport_blueprint,
+        item,
+        item_response,
+        selection_update_behavior,
+        None,
+    );
+}
+
+/// Display a context menu for the provided [`Item`].
+///
+/// The provided `local_data` item will be passed to the context menu actions.
+pub fn context_menu_ui_for_item_with_local_data<T: 'static>(
+    ctx: &ViewerContext<'_>,
+    viewport_blueprint: &ViewportBlueprint,
+    item: &Item,
+    item_response: &egui::Response,
+    selection_update_behavior: SelectionUpdateBehavior,
+    local_data: T,
+) {
+    let mut local_data_map = type_map::TypeMap::new();
+    local_data_map.insert(local_data);
+
+    context_menu_ui_for_item_with_local_data_map(
+        ctx,
+        viewport_blueprint,
+        item,
+        item_response,
+        selection_update_behavior,
+        Some(&local_data_map),
+    );
+}
+
+/// Display a context menu for the provided [`Item`]
+///
+/// The provided `local_data` [`type_map::TypeMap`] will be passed to the context menu actions.
+pub fn context_menu_ui_for_item_with_local_data_map(
+    ctx: &ViewerContext<'_>,
+    viewport_blueprint: &ViewportBlueprint,
+    item: &Item,
+    item_response: &egui::Response,
+    selection_update_behavior: SelectionUpdateBehavior,
+    local_data: Option<&type_map::TypeMap>,
 ) {
     item_response.context_menu(|ui| {
         if ui.input_mut(|i| i.consume_key(egui::Modifiers::NONE, egui::Key::Escape)) {
@@ -55,6 +104,7 @@ pub fn context_menu_ui_for_item(
                 egui_context: ui.ctx().clone(),
                 selection,
                 clicked_item: item,
+                local_data,
             };
             show_context_menu_for_selection(&context_menu_ctx, ui);
         };
@@ -204,6 +254,12 @@ struct ContextMenuContext<'a> {
     egui_context: egui::Context,
     selection: &'a ItemCollection,
     clicked_item: &'a Item,
+
+    /// Custom data provided by the client code to context menu actions.
+    ///
+    /// Action may use this data (if any) to decide if they are locally supported and/or to affect
+    /// their behavior when triggered.
+    local_data: Option<&'a type_map::TypeMap>,
 }
 
 impl<'a> ContextMenuContext<'a> {
@@ -234,6 +290,11 @@ impl<'a> ContextMenuContext<'a> {
                     .map(|container| (container, pos))
             })
     }
+
+    /// Get the local data provided by the client code, if any.
+    pub fn local_date<T: 'static>(&self) -> Option<&T> {
+        self.local_data.and_then(|data| data.get::<T>())
+    }
 }
 
 /// Context menu actions must implement this trait.

crates/viewer/re_time_panel/src/lib.rs

@@ -19,7 +19,7 @@ use std::sync::Arc;
 use egui::emath::Rangef;
 use egui::{pos2, Color32, CursorIcon, NumExt, Painter, PointerButton, Rect, Shape, Ui, Vec2};
 
-use re_context_menu::{context_menu_ui_for_item, SelectionUpdateBehavior};
+use re_context_menu::{context_menu_ui_for_item_with_local_data, SelectionUpdateBehavior};
 use re_data_ui::DataUi as _;
 use re_data_ui::{item_ui::guess_instance_path_icon, sorted_component_list_for_ui};
 use re_entity_db::{EntityDb, EntityTree, InstancePath};
@@ -654,19 +654,14 @@ impl TimePanel {
             .and_then(|item| item.entity_path());
 
         if focused_entity_path.is_some_and(|entity_path| entity_path.is_descendant_of(&tree.path)) {
-            CollapseScope::StreamsTree
+            self.collapse_scope()
                 .entity(tree.path.clone())
                 .set_open(ui.ctx(), true);
         }
 
         // Globally unique id - should only be one of these in view at one time.
         // We do this so that we can support "collapse/expand all" command.
-        let id = egui::Id::new(match self.source {
-            TimePanelSource::Recording => CollapseScope::StreamsTree.entity(tree.path.clone()),
-            TimePanelSource::Blueprint => {
-                CollapseScope::BlueprintStreamsTree.entity(tree.path.clone())
-            }
-        });
+        let id = self.collapse_scope().entity(tree.path.clone()).egui_id();
 
         let list_item::ShowCollapsingResponse {
             item_response: response,
@@ -722,12 +717,13 @@ impl TimePanel {
             }
         }
 
-        context_menu_ui_for_item(
+        context_menu_ui_for_item_with_local_data(
             ctx,
             viewport_blueprint,
             &item.to_item(),
             &response,
             SelectionUpdateBehavior::UseSelection,
+            self.collapse_scope(),
         );
         ctx.handle_select_hover_drag_interactions(&response, item.to_item(), true);
 
@@ -850,12 +846,13 @@ impl TimePanel {
                             .truncate(false),
                     );
 
-                context_menu_ui_for_item(
+                context_menu_ui_for_item_with_local_data(
                     ctx,
                     viewport_blueprint,
                     &item.to_item(),
                     &response,
                     SelectionUpdateBehavior::UseSelection,
+                    self.collapse_scope(),
                 );
                 ctx.handle_select_hover_drag_interactions(&response, item.to_item(), false);
 
@@ -1002,6 +999,13 @@ impl TimePanel {
             });
         }
     }
+
+    fn collapse_scope(&self) -> CollapseScope {
+        match self.source {
+            TimePanelSource::Recording => CollapseScope::StreamsTree,
+            TimePanelSource::Blueprint => CollapseScope::BlueprintStreamsTree,
+        }
+    }
 }
 
 /// Draw the hovered/selected highlight background for a timeline row.