diff --git a/Cargo.lock b/Cargo.lock index d02046f7e1d7..d6bd09fe89e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5782,6 +5782,7 @@ dependencies = [ "re_viewer_context", "re_viewport_blueprint", "static_assertions", + "type-map", ] [[package]] diff --git a/crates/viewer/re_context_menu/Cargo.toml b/crates/viewer/re_context_menu/Cargo.toml index f21a28866eb6..81d96142775c 100644 --- a/crates/viewer/re_context_menu/Cargo.toml +++ b/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 diff --git a/crates/viewer/re_context_menu/src/actions/collapse_expand_all.rs b/crates/viewer/re_context_menu/src/actions/collapse_expand_all.rs index 6d1c583bdf27..0c0408ae3828 100644 --- a/crates/viewer/re_context_menu/src/actions/collapse_expand_all.rs +++ b/crates/viewer/re_context_menu/src/actions/collapse_expand_all.rs @@ -49,9 +49,14 @@ 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), @@ -59,7 +64,9 @@ impl ContextMenuAction for CollapseExpandAllAction { } 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) +} diff --git a/crates/viewer/re_context_menu/src/lib.rs b/crates/viewer/re_context_menu/src/lib.rs index e4c03202c1df..9de623b19954 100644 --- a/crates/viewer/re_context_menu/src/lib.rs +++ b/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( + 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(&self) -> Option<&T> { + self.local_data.and_then(|data| data.get::()) + } } /// Context menu actions must implement this trait. diff --git a/crates/viewer/re_time_panel/src/lib.rs b/crates/viewer/re_time_panel/src/lib.rs index 546d9ae9146b..da2e9d1cd536 100644 --- a/crates/viewer/re_time_panel/src/lib.rs +++ b/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.