Fix group selection/hover with data queries

crates/re_viewer_context/src/query_context.rs

@@ -150,6 +150,17 @@ impl DataResultTree {
         }
     }
 
+    /// Depth-first traversal of a subtree, starting with the given group entity-path, calling `visitor` on each result.
+    pub fn visit_group(
+        &self,
+        entity_path: &EntityPath,
+        visitor: &mut impl FnMut(DataResultHandle),
+    ) {
+        if let Some(subtree_handle) = self.data_results_by_path.get(&(entity_path.clone(), true)) {
+            self.visit_recursive(*subtree_handle, visitor);
+        }
+    }
+
     /// Look up a [`DataResult`] in the tree based on its handle.
     pub fn lookup_result(&self, handle: DataResultHandle) -> Option<&DataResult> {
         self.data_results.get(handle).map(|node| &node.data_result)

crates/re_viewport/src/space_view.rs

@@ -251,7 +251,6 @@ impl SpaceViewBlueprint {
 
         let query_result = ctx.lookup_query_result(self.query_id()).clone();
 
-        // TODO(#4377): Use PerSystemDataResults
         let mut per_system_entities = PerSystemEntities::default();
         {
             re_tracing::profile_scope!("per_system_data_results");

crates/re_viewport/src/space_view_heuristics.rs

@@ -80,8 +80,6 @@ pub fn all_possible_space_views(
                     let mut entity_path_filter = EntityPathFilter::default();
                     entity_path_filter.add_subtree(candidate_space_path.clone());
 
-                    // TODO(#4377): The need to run a query-per-candidate for all possible candidates
-                    // is way too expensive. This needs to be optimized significantly.
                     let candidate_query =
                         DataQueryBlueprint::new(class_identifier, entity_path_filter);
 
@@ -217,7 +215,6 @@ pub fn default_created_space_views(
     // Main pass through all candidates.
     // We first check if a candidate is "interesting" and then split it up/modify it further if required.
     for (candidate, query_result) in candidates {
-        // TODO(#4377): Can spawn heuristics consume the query_result directly?
         let mut per_system_entities = PerSystemEntities::default();
         {
             re_tracing::profile_scope!("per_system_data_results");

crates/re_viewport/src/space_view_highlights.rs

@@ -1,21 +1,16 @@
-use std::collections::BTreeMap;
-
 use egui::NumExt;
 use nohash_hasher::IntMap;
 
 use re_log_types::EntityPathHash;
 use re_renderer::OutlineMaskPreference;
 use re_viewer_context::{
-    ApplicationSelectionState, HoverHighlight, Item, SelectionHighlight, SpaceViewEntityHighlight,
-    SpaceViewHighlights, SpaceViewId, SpaceViewOutlineMasks,
+    HoverHighlight, Item, SelectionHighlight, SpaceViewEntityHighlight, SpaceViewHighlights,
+    SpaceViewId, SpaceViewOutlineMasks,
 };
 
-use crate::SpaceViewBlueprint;
-
 pub fn highlights_for_space_view(
-    selection_state: &ApplicationSelectionState,
+    ctx: &re_viewer_context::ViewerContext<'_>,
     space_view_id: SpaceViewId,
-    _space_views: &BTreeMap<SpaceViewId, SpaceViewBlueprint>,
 ) -> SpaceViewHighlights {
     re_tracing::profile_function!();
 
@@ -36,36 +31,36 @@ pub fn highlights_for_space_view(
         OutlineMaskPreference::some(hover_mask_index, 0)
     };
 
-    for current_selection in selection_state.current().iter_items() {
+    for current_selection in ctx.selection_state().current().iter_items() {
         match current_selection {
             Item::ComponentPath(_) | Item::SpaceView(_) | Item::Container(_) => {}
 
-            Item::DataBlueprintGroup(_space_view_id, _query_id, _entity_path) => {
-                // TODO(#4377): Fix DataBlueprintGroup
-                /*
+            Item::DataBlueprintGroup(group_space_view_id, query_id, group_entity_path) => {
+                // Unlike for selected objects/data we are more picky for data blueprints with our hover highlights
+                // since they are truly local to a space view.
                 if *group_space_view_id == space_view_id {
-                    if let Some(space_view) = space_views.get(group_space_view_id) {
-                        // Everything in the same group should receive the same selection outline.
-                        // (Due to the way outline masks work in re_renderer, we can't leave the hover channel empty)
-                        let selection_mask = next_selection_mask();
-
-                        space_view.contents.visit_group_entities_recursively(
-                            *group_handle,
-                            &mut |entity_path: &EntityPath| {
+                    // Everything in the same group should receive the same selection outline.
+                    // (Due to the way outline masks work in re_renderer, we can't leave the hover channel empty)
+                    let selection_mask = next_selection_mask();
+
+                    let query_result = ctx.lookup_query_result(*query_id).clone();
+
+                    query_result
+                        .tree
+                        .visit_group(group_entity_path, &mut |handle| {
+                            if let Some(result) = query_result.tree.lookup_result(handle) {
                                 highlighted_entity_paths
-                                    .entry(entity_path.hash())
+                                    .entry(result.entity_path.hash())
                                     .or_default()
                                     .overall
                                     .selection = SelectionHighlight::SiblingSelection;
                                 let outline_mask_ids =
-                                    outlines_masks.entry(entity_path.hash()).or_default();
+                                    outlines_masks.entry(result.entity_path.hash()).or_default();
                                 outline_mask_ids.overall =
                                     selection_mask.with_fallback_to(outline_mask_ids.overall);
-                            },
-                        );
-                    }
+                            }
+                        });
                 }
-                */
             }
 
             Item::InstancePath(selected_space_view_context, selected_instance) => {
@@ -113,35 +108,34 @@ pub fn highlights_for_space_view(
         };
     }
 
-    for current_hover in selection_state.hovered().iter_items() {
+    for current_hover in ctx.selection_state().hovered().iter_items() {
         match current_hover {
             Item::ComponentPath(_) | Item::SpaceView(_) | Item::Container(_) => {}
 
-            Item::DataBlueprintGroup(_space_view_id, _query_id, _entity_path) => {
-                // TODO(#4377): Fix DataBlueprintGroup
-                /*
+            Item::DataBlueprintGroup(group_space_view_id, query_id, group_entity_path) => {
                 // Unlike for selected objects/data we are more picky for data blueprints with our hover highlights
                 // since they are truly local to a space view.
                 if *group_space_view_id == space_view_id {
-                    if let Some(space_view) = space_views.get(group_space_view_id) {
-                        // Everything in the same group should receive the same selection outline.
-                        let hover_mask = next_hover_mask();
+                    // Everything in the same group should receive the same hover outline.
+                    let hover_mask = next_hover_mask();
+
+                    let query_result = ctx.lookup_query_result(*query_id).clone();
 
-                        space_view.contents.visit_group_entities_recursively(
-                            *group_handle,
-                            &mut |entity_path: &EntityPath| {
+                    query_result
+                        .tree
+                        .visit_group(group_entity_path, &mut |handle| {
+                            if let Some(result) = query_result.tree.lookup_result(handle) {
                                 highlighted_entity_paths
-                                    .entry(entity_path.hash())
+                                    .entry(result.entity_path.hash())
                                     .or_default()
                                     .overall
                                     .hover = HoverHighlight::Hovered;
-                                let mask = outlines_masks.entry(entity_path.hash()).or_default();
+                                let mask =
+                                    outlines_masks.entry(result.entity_path.hash()).or_default();
                                 mask.overall = hover_mask.with_fallback_to(mask.overall);
-                            },
-                        );
-                    }
+                            }
+                        });
                 }
-                */
             }
 
             Item::InstancePath(_, selected_instance) => {

crates/re_viewport/src/system_execution.rs

@@ -73,8 +73,7 @@ pub fn execute_systems_for_space_views<'a>(
         .par_drain(..)
         .filter_map(|space_view_id| {
             let space_view_blueprint = space_views.get(&space_view_id)?;
-            let highlights =
-                highlights_for_space_view(ctx.selection_state(), space_view_id, space_views);
+            let highlights = highlights_for_space_view(ctx, space_view_id);
             let output = space_view_blueprint.execute_systems(ctx, time_int, highlights);
             Some((space_view_id, output))
         })

crates/re_viewport/src/viewport.rs

@@ -495,11 +495,7 @@ impl<'a, 'b> egui_tiles::Behavior<SpaceViewId> for TabViewer<'a, 'b> {
             if let Some(result) = self.executed_systems_per_space_view.remove(space_view_id) {
                 result
             } else {
-                let highlights = highlights_for_space_view(
-                    self.ctx.selection_state(),
-                    *space_view_id,
-                    self.space_views,
-                );
+                let highlights = highlights_for_space_view(self.ctx, *space_view_id);
                 space_view_blueprint.execute_systems(self.ctx, latest_at, highlights)
             };