Improve performance for many entities

crates/re_space_view/src/lib.rs

@@ -3,10 +3,10 @@
 //! Types & utilities for defining Space View classes and communicating with the Viewport.
 
 pub mod controls;
-mod data_blueprint;
 mod screenshot;
+mod space_view_contents;
 mod unreachable_transform_reason;
 
-pub use data_blueprint::{DataBlueprintGroup, DataBlueprintTree};
 pub use screenshot::ScreenshotMode;
+pub use space_view_contents::{DataBlueprintGroup, SpaceViewContents};
 pub use unreachable_transform_reason::UnreachableTransformReason;

crates/re_space_view/src/data_blueprint.rs → crates/re_space_view/src/space_view_contents.rs

@@ -1,8 +1,8 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 
-use nohash_hasher::{IntMap, IntSet};
+use nohash_hasher::IntMap;
 use re_data_store::{EntityPath, EntityProperties, EntityPropertyMap};
-use re_viewer_context::DataBlueprintGroupHandle;
+use re_viewer_context::{DataBlueprintGroupHandle, PerSystemEntities};
 use slotmap::SlotMap;
 use smallvec::{smallvec, SmallVec};
 
@@ -79,7 +79,7 @@ impl DataBlueprints {
 
 /// Tree of all data blueprint groups for a single space view.
 #[derive(Clone, serde::Deserialize, serde::Serialize)]
-pub struct DataBlueprintTree {
+pub struct SpaceViewContents {
     /// All data blueprint groups.
     groups: SlotMap<DataBlueprintGroupHandle, DataBlueprintGroup>,
 
@@ -92,11 +92,12 @@ pub struct DataBlueprintTree {
 
     /// List of all entities that we query via this data blueprint collection.
     ///
-    /// Two things to keep in sync:
-    /// * children on [`DataBlueprintGroup`] this is on
-    /// * elements in [`Self::path_to_group`]
-    /// TODO(andreas): Can we reduce the amount of these dependencies?
-    entity_paths: IntSet<EntityPath>,
+    /// Currently this is reset every frame in `SpaceViewBlueprint::reset_systems_per_entity_path`.
+    /// In the future, we may want to keep this around and only add/remove systems
+    /// for entities. But at this point we'd likely handle the heuristics a bit differently as well
+    /// and don't use serde here for serialization.
+    #[serde(skip)]
+    per_system_entity_list: PerSystemEntities,
 
     /// Root group, always exists as a placeholder
     root_group_handle: DataBlueprintGroupHandle,
@@ -105,12 +106,12 @@ pub struct DataBlueprintTree {
 }
 
 /// Determine whether this `DataBlueprintTree` has user-edits relative to another `DataBlueprintTree`
-impl DataBlueprintTree {
+impl SpaceViewContents {
     pub fn has_edits(&self, other: &Self) -> bool {
         let Self {
             groups,
             path_to_group,
-            entity_paths,
+            per_system_entity_list: _,
             root_group_handle,
             data_blueprints,
         } = self;
@@ -123,13 +124,12 @@ impl DataBlueprintTree {
                     .map_or(true, |other_val| val.has_edits(other_val))
             })
             || *path_to_group != other.path_to_group
-            || *entity_paths != other.entity_paths
             || *root_group_handle != other.root_group_handle
             || data_blueprints.has_edits(&other.data_blueprints)
     }
 }
 
-impl Default for DataBlueprintTree {
+impl Default for SpaceViewContents {
     fn default() -> Self {
         let mut groups = SlotMap::default();
         let root_group = groups.insert(DataBlueprintGroup::default());
@@ -140,14 +140,14 @@ impl Default for DataBlueprintTree {
         Self {
             groups,
             path_to_group: path_to_blueprint,
-            entity_paths: IntSet::default(),
+            per_system_entity_list: BTreeMap::default(),
             root_group_handle: root_group,
             data_blueprints: DataBlueprints::default(),
         }
     }
 }
 
-impl DataBlueprintTree {
+impl SpaceViewContents {
     /// Returns a handle to the root data blueprint.
     ///
     /// Even if there are no other groups, we always have a root group at the top.
@@ -207,12 +207,47 @@ impl DataBlueprintTree {
     }
 
     pub fn contains_entity(&self, path: &EntityPath) -> bool {
-        self.entity_paths.contains(path)
+        // If an entity is in path_to_group it is *likely* also an entity in the Space View.
+        // However, it could be that the path *only* refers to a group, not also an entity.
+        // So once we resolved the group, we need to check if it contains the entity of interest.
+        self.path_to_group
+            .get(path)
+            .and_then(|group| {
+                self.groups
+                    .get(*group)
+                    .and_then(|group| group.entities.get(path))
+            })
+            .is_some()
     }
 
     /// List of all entities that we query via this data blueprint collection.
-    pub fn entity_paths(&self) -> &IntSet<EntityPath> {
-        &self.entity_paths
+    pub fn entity_paths(&self) -> impl Iterator<Item = &EntityPath> {
+        // Each entity is only ever in one group, therefore collecting all entities from all groups, gives us all entities.
+        self.groups.values().flat_map(|group| group.entities.iter())
+    }
+
+    pub fn per_system_entities(&self) -> &PerSystemEntities {
+        &self.per_system_entity_list
+    }
+
+    pub fn per_system_entities_mut(&mut self) -> &mut PerSystemEntities {
+        &mut self.per_system_entity_list
+    }
+
+    pub fn contains_all_entities_from(&self, other: &SpaceViewContents) -> bool {
+        for (system, entities) in &other.per_system_entity_list {
+            let Some(self_entities) = self.per_system_entity_list.get(system) else {
+                if entities.is_empty() {
+                    continue;
+                } else {
+                    return false;
+                }
+            };
+            if !entities.is_subset(self_entities) {
+                return false;
+            }
+        }
+        true
     }
 
     /// Should be called on frame start.
@@ -225,7 +260,7 @@ impl DataBlueprintTree {
         // and/or when new entity paths are added, but such memoization would add complexity.
 
         fn project_tree(
-            tree: &mut DataBlueprintTree,
+            tree: &mut SpaceViewContents,
             parent_properties: &EntityProperties,
             group_handle: DataBlueprintGroupHandle,
         ) {
@@ -270,8 +305,6 @@ impl DataBlueprintTree {
         let mut new_leaf_groups = Vec::new();
 
         for path in paths {
-            self.entity_paths.insert(path.clone());
-
             // Is there already a group associated with this exact path? (maybe because a child was logged there earlier)
             // If so, we can simply move it under this existing group.
             let group_handle = if let Some(group_handle) = self.path_to_group.get(path) {
@@ -393,14 +426,19 @@ impl DataBlueprintTree {
     ///
     /// If the entity was not known by this data blueprint tree nothing happens.
     pub fn remove_entity(&mut self, path: &EntityPath) {
+        re_tracing::profile_function!();
+
         if let Some(group_handle) = self.path_to_group.get(path) {
             if let Some(group) = self.groups.get_mut(*group_handle) {
                 group.entities.remove(path);
                 self.remove_group_if_empty(*group_handle);
             }
         }
         self.path_to_group.remove(path);
-        self.entity_paths.remove(path);
+
+        for per_system_list in self.per_system_entity_list.values_mut() {
+            per_system_list.remove(path);
+        }
     }
 
     /// Removes a group and all its entities and subgroups from the blueprint tree
@@ -421,7 +459,9 @@ impl DataBlueprintTree {
 
         // Remove all child entities.
         for entity_path in &group.entities {
-            self.entity_paths.remove(entity_path);
+            for per_system_list in self.per_system_entity_list.values_mut() {
+                per_system_list.remove(entity_path);
+            }
         }
 
         // Remove from `path_to_group` map.

crates/re_space_view_bar_chart/src/view_part_system.rs

@@ -6,8 +6,8 @@ use re_data_store::EntityPath;
 use re_log_types::{TimeInt, Timeline};
 use re_types::{ComponentName, Loggable as _};
 use re_viewer_context::{
-    ArchetypeDefinition, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem,
-    ViewQuery, ViewerContext,
+    ArchetypeDefinition, NamedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection,
+    ViewPartSystem, ViewQuery, ViewerContext,
 };
 
 /// A bar chart system, with everything needed to render it.
@@ -16,6 +16,12 @@ pub struct BarChartViewPartSystem {
     pub charts: BTreeMap<EntityPath, Tensor>,
 }
 
+impl NamedViewSystem for BarChartViewPartSystem {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "BarChartView".into()
+    }
+}
+
 impl ViewPartSystem for BarChartViewPartSystem {
     fn archetype(&self) -> ArchetypeDefinition {
         vec1::vec1![Tensor::name()]
@@ -51,7 +57,7 @@ impl ViewPartSystem for BarChartViewPartSystem {
 
         let store = &ctx.store_db.entity_db.data_store;
 
-        for (ent_path, props) in query.iter_entities() {
+        for (ent_path, props) in query.iter_entities_for_system(Self::name()) {
             if !props.visible {
                 continue;
             }

crates/re_space_view_spatial/src/contexts/annotation_context.rs

@@ -1,9 +1,17 @@
 use re_types::{components::AnnotationContext, Loggable};
-use re_viewer_context::{AnnotationMap, ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{
+    AnnotationMap, ArchetypeDefinition, NamedViewSystem, ViewContextSystem, ViewSystemName,
+};
 
 #[derive(Default)]
 pub struct AnnotationSceneContext(pub AnnotationMap);
 
+impl NamedViewSystem for AnnotationSceneContext {
+    fn name() -> ViewSystemName {
+        "AnnotationSceneContext".into()
+    }
+}
+
 impl ViewContextSystem for AnnotationSceneContext {
     fn archetypes(&self) -> Vec<ArchetypeDefinition> {
         vec![vec1::vec1![AnnotationContext::name()]]
@@ -15,11 +23,10 @@ impl ViewContextSystem for AnnotationSceneContext {
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
         re_tracing::profile_function!();
-        self.0.load(
-            ctx,
-            &query.latest_at_query(),
-            query.iter_entities().map(|(p, _)| p),
-        );
+        // We create a list of *all* entities here, do not only iterate over those with annotation context.
+        // TODO(andreas): But knowing ahead of time where we have annotation contexts could be used for optimization.
+        self.0
+            .load(ctx, &query.latest_at_query(), query.iter_entities());
     }
 
     fn as_any(&self) -> &dyn std::any::Any {

crates/re_space_view_spatial/src/contexts/depth_offsets.rs

@@ -4,7 +4,7 @@ use nohash_hasher::IntMap;
 
 use re_log_types::EntityPathHash;
 use re_types::{components::DrawOrder, Loggable as _};
-use re_viewer_context::{ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{ArchetypeDefinition, NamedViewSystem, ViewContextSystem};
 
 /// Context for creating a mapping from [`DrawOrder`] to [`re_renderer::DepthOffset`].
 #[derive(Default)]
@@ -19,6 +19,12 @@ pub struct EntityDepthOffsets {
     pub points: re_renderer::DepthOffset,
 }
 
+impl NamedViewSystem for EntityDepthOffsets {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "EntityDepthOffsets".into()
+    }
+}
+
 impl ViewContextSystem for EntityDepthOffsets {
     fn archetypes(&self) -> Vec<ArchetypeDefinition> {
         vec![vec1::vec1![DrawOrder::name()]]
@@ -29,8 +35,6 @@ impl ViewContextSystem for EntityDepthOffsets {
         ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
-        re_tracing::profile_function!();
-
         #[derive(PartialEq, PartialOrd, Eq, Ord)]
         enum DrawOrderTarget {
             Entity(EntityPathHash),
@@ -44,7 +48,7 @@ impl ViewContextSystem for EntityDepthOffsets {
 
         // Use a BTreeSet for entity hashes to get a stable order.
         let mut entities_per_draw_order = BTreeMap::<DrawOrder, BTreeSet<DrawOrderTarget>>::new();
-        for (ent_path, _) in query.iter_entities() {
+        for (ent_path, _) in query.iter_entities_for_system(Self::name()) {
             if let Some(draw_order) = store.query_latest_component::<DrawOrder>(
                 ent_path,
                 &ctx.rec_cfg.time_ctrl.current_query(),

crates/re_space_view_spatial/src/contexts/mod.rs

@@ -15,7 +15,9 @@ pub use transform_context::{pinhole_camera_view_coordinates, TransformContext};
 // -----------------------------------------------------------------------------
 
 use re_renderer::DepthOffset;
-use re_viewer_context::{Annotations, SpaceViewClassRegistryError, ViewContextSystem};
+use re_viewer_context::{
+    Annotations, NamedViewSystem, SpaceViewClassRegistryError, ViewContextSystem,
+};
 
 /// Context objects for a single entity in a spatial scene.
 pub struct SpatialSceneEntityContext<'a> {
@@ -32,6 +34,12 @@ pub struct PrimitiveCounter {
     pub num_primitives: AtomicUsize,
 }
 
+impl NamedViewSystem for PrimitiveCounter {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "PrimitiveCounter".into()
+    }
+}
+
 impl ViewContextSystem for PrimitiveCounter {
     fn archetypes(&self) -> Vec<re_viewer_context::ArchetypeDefinition> {
         Vec::new()

crates/re_space_view_spatial/src/contexts/non_interactive_entities.rs

@@ -1,13 +1,19 @@
 use nohash_hasher::IntSet;
 use re_log_types::EntityPathHash;
-use re_viewer_context::ViewContextSystem;
+use re_viewer_context::{NamedViewSystem, ViewContextSystem};
 
 /// List of all non-interactive entities for lookup during picking evaluation.
 ///
 /// TODO(wumpf/jleibs): This is a temporary solution until the picking code can query propagated blueprint properties directly.
 #[derive(Default)]
 pub struct NonInteractiveEntities(pub IntSet<EntityPathHash>);
 
+impl NamedViewSystem for NonInteractiveEntities {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "NonInteractiveEntities".into()
+    }
+}
+
 impl ViewContextSystem for NonInteractiveEntities {
     fn archetypes(&self) -> Vec<re_viewer_context::ArchetypeDefinition> {
         Vec::new()
@@ -18,8 +24,6 @@ impl ViewContextSystem for NonInteractiveEntities {
         _ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
     ) {
-        re_tracing::profile_function!();
-
         self.0 = query
             .entity_props_map
             .iter()

crates/re_space_view_spatial/src/contexts/shared_render_builders.rs

@@ -1,6 +1,6 @@
 use parking_lot::{MappedMutexGuard, Mutex, MutexGuard};
 use re_renderer::{LineStripSeriesBuilder, PointCloudBuilder, RenderContext};
-use re_viewer_context::{ArchetypeDefinition, ViewContextSystem};
+use re_viewer_context::{ArchetypeDefinition, NamedViewSystem, ViewContextSystem};
 
 use crate::parts::{
     SIZE_BOOST_IN_POINTS_FOR_LINE_OUTLINES, SIZE_BOOST_IN_POINTS_FOR_POINT_OUTLINES,
@@ -14,6 +14,12 @@ pub struct SharedRenderBuilders {
     pub points: Mutex<Option<PointCloudBuilder>>,
 }
 
+impl NamedViewSystem for SharedRenderBuilders {
+    fn name() -> re_viewer_context::ViewSystemName {
+        "SharedRenderBuilders".into()
+    }
+}
+
 impl SharedRenderBuilders {
     pub fn lines(&self) -> MappedMutexGuard<'_, LineStripSeriesBuilder> {
         MutexGuard::map(self.lines.lock(), |l| l.as_mut().unwrap())