Fix clicking object with single instance (of every component) selecti…

…ng instance instead of entity (#2573)

### What

Fixes a regression introduced by #2504

Single click on box in objectron:
Before:

![image](https://github.com/rerun-io/rerun/assets/1220815/2ec91fa9-1120-42b0-bf5f-ba5b4a1f290b)

After:

![image](https://github.com/rerun-io/rerun/assets/1220815/93746403-d22d-4626-97fc-642a550aa36c)


tested with `python ./scripts/run_all.py`, in particular checked on
selection history and point clouds behaving as they should via #2504

### 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 [demo.rerun.io](https://demo.rerun.io/pr/2573) (if
applicable)

- [PR Build Summary](https://build.rerun.io/pr/2573)
- [Docs
preview](https://rerun.io/preview/pr%3Aandreas%2Ffix-single-instance-select/docs)
- [Examples
preview](https://rerun.io/preview/pr%3Aandreas%2Ffix-single-instance-select/examples)

crates/re_data_ui/src/item_ui.rs

@@ -5,8 +5,7 @@
 use re_data_store::InstancePath;
 use re_log_types::{ComponentPath, EntityPath, TimeInt, Timeline};
 use re_viewer_context::{
-    DataBlueprintGroupHandle, HoverHighlight, Item, SelectionState, SpaceViewId, UiVerbosity,
-    ViewerContext,
+    DataBlueprintGroupHandle, HoverHighlight, Item, SpaceViewId, UiVerbosity, ViewerContext,
 };
 
 use super::DataUi;
@@ -109,7 +108,7 @@ pub fn instance_path_button_to(
             instance_path.data_ui(ctx, ui, UiVerbosity::Reduced, &ctx.current_query());
         });
 
-    cursor_interact_with_selectable(ctx.selection_state_mut(), response, item)
+    cursor_interact_with_selectable(ctx, response, item)
 }
 
 /// Show a component path and make it selectable.
@@ -135,7 +134,7 @@ pub fn component_path_button_to(
 ) -> egui::Response {
     let item = Item::ComponentPath(component_path.clone());
     let response = ui.selectable_label(ctx.selection().contains(&item), text);
-    cursor_interact_with_selectable(ctx.selection_state_mut(), response, item)
+    cursor_interact_with_selectable(ctx, response, item)
 }
 
 pub fn data_blueprint_group_button_to(
@@ -155,7 +154,7 @@ pub fn data_blueprint_group_button_to(
             ctx.selection().contains(&item),
         )
         .on_hover_text("Group");
-    cursor_interact_with_selectable(ctx.selection_state_mut(), response, item)
+    cursor_interact_with_selectable(ctx, response, item)
 }
 
 pub fn data_blueprint_button_to(
@@ -176,7 +175,7 @@ pub fn data_blueprint_button_to(
             ui.label(format!("Path: {entity_path}"));
             entity_path.data_ui(ctx, ui, UiVerbosity::Reduced, &ctx.current_query());
         });
-    cursor_interact_with_selectable(ctx.selection_state_mut(), response, item)
+    cursor_interact_with_selectable(ctx, response, item)
 }
 
 pub fn time_button(
@@ -225,14 +224,14 @@ pub fn timeline_button_to(
 
 // TODO(andreas): Move elsewhere, this is not directly part of the item_ui.
 pub fn cursor_interact_with_selectable(
-    selection_state: &mut SelectionState,
+    ctx: &mut ViewerContext<'_>,
     response: egui::Response,
     item: Item,
 ) -> egui::Response {
     let is_item_hovered =
-        selection_state.highlight_for_ui_element(&item) == HoverHighlight::Hovered;
+        ctx.selection_state().highlight_for_ui_element(&item) == HoverHighlight::Hovered;
 
-    select_hovered_on_click(&response, selection_state, &[item]);
+    select_hovered_on_click(ctx, &response, &[item]);
     // TODO(andreas): How to deal with shift click for selecting ranges?
 
     if is_item_hovered {
@@ -244,19 +243,22 @@ pub fn cursor_interact_with_selectable(
 
 // TODO(andreas): Move elsewhere, this is not directly part of the item_ui.
 pub fn select_hovered_on_click(
+    ctx: &mut ViewerContext<'_>,
     response: &egui::Response,
-    selection_state: &mut SelectionState,
     items: &[Item],
 ) {
+    re_tracing::profile_function!();
+
     if response.hovered() {
-        selection_state.set_hovered(items.iter().cloned());
+        ctx.selection_state_mut().set_hovered(items.iter().cloned());
     }
 
     if response.clicked() {
         if response.ctx.input(|i| i.modifiers.command) {
-            selection_state.toggle_selection(items.to_vec());
+            ctx.selection_state_mut().toggle_selection(items.to_vec());
         } else {
-            selection_state.set_selection(items.iter().cloned());
+            ctx.selection_state_mut()
+                .set_selection(items.iter().cloned());
         }
     }
 }

crates/re_space_view_spatial/src/ui.rs

@@ -11,8 +11,9 @@ use re_format::format_f32;
 use re_renderer::OutlineConfig;
 use re_space_view::ScreenshotMode;
 use re_viewer_context::{
-    HoverHighlight, HoveredSpace, Item, SelectionHighlight, SpaceViewHighlights, SpaceViewId,
-    SpaceViewState, TensorDecodeCache, TensorStatsCache, UiVerbosity, ViewerContext,
+    resolve_mono_instance_path, HoverHighlight, HoveredSpace, Item, SelectionHighlight,
+    SpaceViewHighlights, SpaceViewId, SpaceViewState, TensorDecodeCache, TensorStatsCache,
+    UiVerbosity, ViewerContext,
 };
 
 use super::{
@@ -731,8 +732,8 @@ pub fn picking(
         let picked_image_with_coords = if hit.hit_type == PickingHitType::TexturedRect
             || is_depth_cloud
         {
-            let store = &ctx.store_db.entity_db.data_store;
-            store
+            ctx.store_db
+                .store()
                 .query_latest_component::<Tensor>(&instance_path.entity_path, &ctx.current_query())
                 .and_then(|tensor| {
                     // If we're here because of back-projection, but this wasn't actually a depth image, drop out.
@@ -758,6 +759,12 @@ pub fn picking(
         if picked_image_with_coords.is_some() {
             // We don't support selecting pixels yet.
             instance_path.instance_key = re_log_types::InstanceKey::SPLAT;
+        } else {
+            instance_path = resolve_mono_instance_path(
+                &ctx.current_query(),
+                ctx.store_db.store(),
+                &instance_path,
+            );
         }
 
         hovered_items.push(Item::InstancePath(
@@ -844,8 +851,7 @@ pub fn picking(
         };
     }
 
-    item_ui::select_hovered_on_click(&response, ctx.selection_state_mut(), &hovered_items);
-    ctx.set_hovered(hovered_items.into_iter());
+    item_ui::select_hovered_on_click(ctx, &response, &hovered_items);
 
     let hovered_space = match state.nav_mode.get() {
         SpatialNavigationMode::TwoD => HoveredSpace::TwoD {

crates/re_time_panel/src/data_density_graph.rs

@@ -370,7 +370,7 @@ pub fn data_density_graph_ui(
     time_histogram: &TimeHistogram,
     row_rect: Rect,
     time_ranges_ui: &TimeRangesUi,
-    item: Item,
+    item: &Item,
 ) {
     re_tracing::profile_function!();
 
@@ -470,7 +470,7 @@ pub fn data_density_graph_ui(
         data_dentity_graph_painter,
         row_rect.y_range(),
         time_area_painter,
-        graph_color(ctx, &item, ui),
+        graph_color(ctx, item, ui),
         hovered_x_range,
     );
 
@@ -487,7 +487,7 @@ pub fn data_density_graph_ui(
             show_row_ids_tooltip(
                 ctx,
                 ui.ctx(),
-                &item,
+                item,
                 hovered_time_range,
                 num_hovered_messages,
             );

crates/re_time_panel/src/lib.rs

@@ -448,7 +448,7 @@ impl TimePanel {
                 num_messages_at_time,
                 row_rect,
                 &self.time_ranges_ui,
-                item,
+                &item,
             );
         }
     }
@@ -544,7 +544,7 @@ impl TimePanel {
                         messages_over_time,
                         row_rect,
                         &self.time_ranges_ui,
-                        item,
+                        &item,
                     );
                 }
             }

crates/re_viewer_context/src/item.rs

@@ -120,3 +120,55 @@ impl std::iter::IntoIterator for ItemCollection {
         self.items.into_iter()
     }
 }
+
+/// If the given item refers to the first element of an instance with a single element, resolve to a splatted entity path.
+pub fn resolve_mono_instance_path_item(
+    query: &re_arrow_store::LatestAtQuery,
+    store: &re_arrow_store::DataStore,
+    item: &Item,
+) -> Item {
+    // Resolve to entity path if there's only a single instance.
+    match item {
+        Item::InstancePath(space_view, instance) => Item::InstancePath(
+            *space_view,
+            resolve_mono_instance_path(query, store, instance),
+        ),
+        Item::ComponentPath(_) | Item::SpaceView(_) | Item::DataBlueprintGroup(_, _) => {
+            item.clone()
+        }
+    }
+}
+
+/// If the given path refers to the first element of an instance with a single element, resolve to a splatted entity path.
+pub fn resolve_mono_instance_path(
+    query: &re_arrow_store::LatestAtQuery,
+    store: &re_arrow_store::DataStore,
+    instance: &re_data_store::InstancePath,
+) -> re_data_store::InstancePath {
+    re_tracing::profile_function!();
+
+    if instance.instance_key.0 == 0 {
+        let Some(components) = store
+                .all_components(&query.timeline, &instance.entity_path) else {
+            // No components at all, return splatted entity.
+            return re_data_store::InstancePath::entity_splat(instance.entity_path.clone());
+        };
+        for component in components {
+            if let Some((_row_id, instances)) = re_query::get_component_with_instances(
+                store,
+                query,
+                &instance.entity_path,
+                component,
+            ) {
+                if instances.len() > 1 {
+                    return instance.clone();
+                }
+            }
+        }
+
+        // All instances had only a single element or less, resolve to splatted entity.
+        return re_data_store::InstancePath::entity_splat(instance.entity_path.clone());
+    }
+
+    instance.clone()
+}

crates/re_viewer_context/src/lib.rs

@@ -29,7 +29,7 @@ pub use command_sender::{
     command_channel, CommandReceiver, CommandSender, SystemCommand, SystemCommandSender,
 };
 pub use component_ui_registry::{ComponentUiRegistry, UiVerbosity};
-pub use item::{Item, ItemCollection};
+pub use item::{resolve_mono_instance_path, resolve_mono_instance_path_item, Item, ItemCollection};
 pub use selection_history::SelectionHistory;
 pub use selection_state::{
     HoverHighlight, HoveredSpace, InteractionHighlight, SelectionHighlight, SelectionState,

crates/re_viewer_context/src/viewer_context.rs

@@ -1,8 +1,8 @@
 use re_data_store::store_db::StoreDb;
 
 use crate::{
-    AppOptions, Caches, CommandSender, ComponentUiRegistry, Item, ItemCollection, SelectionState,
-    SpaceViewClassRegistry, StoreContext, TimeControl,
+    item::resolve_mono_instance_path_item, AppOptions, Caches, CommandSender, ComponentUiRegistry,
+    Item, ItemCollection, SelectionState, SpaceViewClassRegistry, StoreContext, TimeControl,
 };
 
 /// Common things needed by many parts of the viewer.
@@ -45,8 +45,14 @@ impl<'a> ViewerContext<'a> {
     /// Sets a single selection, updating history as needed.
     ///
     /// Returns the previous selection.
-    pub fn set_single_selection(&mut self, item: Item) -> ItemCollection {
-        self.rec_cfg.selection_state.set_single_selection(item)
+    pub fn set_single_selection(&mut self, item: &Item) -> ItemCollection {
+        self.rec_cfg
+            .selection_state
+            .set_single_selection(resolve_mono_instance_path_item(
+                &self.rec_cfg.time_ctrl.current_query(),
+                self.store_db.store(),
+                item,
+            ))
     }
 
     /// Returns the current selection.
@@ -60,8 +66,16 @@ impl<'a> ViewerContext<'a> {
     }
 
     /// Set the hovered objects. Will be in [`Self::hovered`] on the next frame.
-    pub fn set_hovered(&mut self, hovered: impl Iterator<Item = Item>) {
-        self.rec_cfg.selection_state.set_hovered(hovered);
+    pub fn set_hovered<'b>(&mut self, hovered: impl Iterator<Item = &'b Item>) {
+        self.rec_cfg
+            .selection_state
+            .set_hovered(hovered.map(|item| {
+                resolve_mono_instance_path_item(
+                    &self.rec_cfg.time_ctrl.current_query(),
+                    self.store_db.store(),
+                    item,
+                )
+            }));
     }
 
     pub fn selection_state(&self) -> &SelectionState {

crates/re_viewport/src/lib.rs

@@ -51,6 +51,6 @@ pub mod item_ui {
                 is_selected,
             )
             .on_hover_text("Space View");
-        item_ui::cursor_interact_with_selectable(ctx.selection_state_mut(), response, item)
+        item_ui::cursor_interact_with_selectable(ctx, response, item)
     }
 }

crates/re_viewport/src/viewport.rs

@@ -262,7 +262,7 @@ impl<'a, 'b> egui_tiles::Behavior<SpaceViewId> for TabViewer<'a, 'b> {
         if button_response.clicked() {
             if let Some(egui_tiles::Tile::Pane(space_view_id)) = tiles.get(tile_id) {
                 self.ctx
-                    .set_single_selection(Item::SpaceView(*space_view_id));
+                    .set_single_selection(&Item::SpaceView(*space_view_id));
             }
         }
     }
@@ -315,7 +315,7 @@ impl<'a, 'b> egui_tiles::Behavior<SpaceViewId> for TabViewer<'a, 'b> {
             {
                 *self.maximized = Some(space_view_id);
                 self.ctx
-                    .set_single_selection(Item::SpaceView(space_view_id));
+                    .set_single_selection(&Item::SpaceView(space_view_id));
             }
         }
 

crates/re_viewport/src/viewport_blueprint_ui.rs

@@ -366,7 +366,7 @@ impl ViewportBlueprint<'_> {
                 {
                     ui.close_menu();
                     let new_space_view_id = self.add_space_view(space_view);
-                    ctx.set_single_selection(Item::SpaceView(new_space_view_id));
+                    ctx.set_single_selection(&Item::SpaceView(new_space_view_id));
                 }
             }
         })

examples/rust/custom_space_view/src/color_coordinates_space_view.rs

@@ -247,8 +247,8 @@ fn color_space_ui(
                 instance.data_ui(ctx, ui, UiVerbosity::Reduced, &ctx.current_query());
             });
             item_ui::select_hovered_on_click(
+                ctx,
                 &interact,
-                ctx.selection_state_mut(),
                 &[Item::InstancePath(Some(space_view_id), instance)],
             );
         }