Introduce store subscriber (VisualizerEntitySubscriber) to maintain list of entities per system incrementally

Cargo.toml

@@ -106,6 +106,7 @@ atty = "0.2"
 backtrace = "0.3"
 bincode = "1.3"
 bitflags = { version = "2.4", features = ["bytemuck"] }
+bit-vec = "0.6"
 blackbox = "0.2.0"
 bytemuck = { version = "1.11", features = ["extern_crate_alloc"] }
 camino = "1.1"

crates/re_viewer_context/Cargo.toml

@@ -30,6 +30,7 @@ re_ui.workspace = true
 
 ahash.workspace = true
 anyhow.workspace = true
+bit-vec.workspace = true
 bytemuck.workspace = true
 egui-wgpu.workspace = true
 egui.workspace = true

crates/re_viewer_context/src/space_view/mod.rs

@@ -15,6 +15,7 @@ mod system_execution_output;
 mod view_context_system;
 mod view_part_system;
 mod view_query;
+mod visualizer_entity_subscriber;
 
 pub use auto_spawn_heuristic::AutoSpawnHeuristic;
 pub use dyn_space_view_class::{

crates/re_viewer_context/src/space_view/space_view_class_registry.rs

@@ -1,11 +1,17 @@
 use ahash::{HashMap, HashSet};
+use nohash_hasher::IntSet;
+use re_arrow_store::DataStore;
+use re_log_types::EntityPath;
 
 use crate::{
     DynSpaceViewClass, IdentifiedViewSystem, SpaceViewClassIdentifier, ViewContextCollection,
     ViewContextSystem, ViewPartCollection, ViewPartSystem, ViewSystemIdentifier,
 };
 
-use super::space_view_class_placeholder::SpaceViewClassPlaceholder;
+use super::{
+    space_view_class_placeholder::SpaceViewClassPlaceholder,
+    visualizer_entity_subscriber::VisualizerEntitySubscriber,
+};
 
 #[derive(Debug, thiserror::Error)]
 #[allow(clippy::enum_variant_names)]
@@ -54,7 +60,7 @@ impl SpaceViewSystemRegistrator<'_> {
             self.registry
                 .context_systems
                 .entry(T::identifier())
-                .or_insert_with(|| SystemTypeRegistryEntry {
+                .or_insert_with(|| ContextSystemTypeRegistryEntry {
                     factory_method: Box::new(|| Box::<T>::default()),
                     used_by: Default::default(),
                 })
@@ -91,9 +97,16 @@ impl SpaceViewSystemRegistrator<'_> {
             self.registry
                 .visualizers
                 .entry(T::identifier())
-                .or_insert_with(|| SystemTypeRegistryEntry {
-                    factory_method: Box::new(|| Box::<T>::default()),
-                    used_by: Default::default(),
+                .or_insert_with(|| {
+                    let entity_subscriber_handle = DataStore::register_subscriber(Box::new(
+                        VisualizerEntitySubscriber::new(&T::default()),
+                    ));
+
+                    VisualizerTypeRegistryEntry {
+                        factory_method: Box::new(|| Box::<T>::default()),
+                        used_by: Default::default(),
+                        entity_subscriber_handle,
+                    }
                 })
                 .used_by
                 .insert(self.identifier);
@@ -110,37 +123,53 @@ impl SpaceViewSystemRegistrator<'_> {
 }
 
 /// Space view class entry in [`SpaceViewClassRegistry`].
-struct SpaceViewClassRegistryEntry {
-    class: Box<dyn DynSpaceViewClass>,
-    context_systems: HashSet<ViewSystemIdentifier>,
-    visualizers: HashSet<ViewSystemIdentifier>,
+pub struct SpaceViewClassRegistryEntry {
+    pub class: Box<dyn DynSpaceViewClass>,
+    pub context_system_ids: HashSet<ViewSystemIdentifier>,
+    pub visualizer_system_ids: HashSet<ViewSystemIdentifier>,
 }
 
 #[allow(clippy::derivable_impls)] // Clippy gets this one wrong.
 impl Default for SpaceViewClassRegistryEntry {
     fn default() -> Self {
         Self {
             class: Box::<SpaceViewClassPlaceholder>::default(),
-            context_systems: Default::default(),
-            visualizers: Default::default(),
+            context_system_ids: Default::default(),
+            visualizer_system_ids: Default::default(),
         }
     }
 }
 
-/// System type entry in [`SpaceViewClassRegistry`].
-struct SystemTypeRegistryEntry<T: ?Sized> {
-    factory_method: Box<dyn Fn() -> Box<T> + Send + Sync>,
+/// Context system type entry in [`SpaceViewClassRegistry`].
+struct ContextSystemTypeRegistryEntry {
+    factory_method: Box<dyn Fn() -> Box<dyn ViewContextSystem> + Send + Sync>,
     used_by: HashSet<SpaceViewClassIdentifier>,
 }
 
+/// Visualizer entry in [`SpaceViewClassRegistry`].
+struct VisualizerTypeRegistryEntry {
+    factory_method: Box<dyn Fn() -> Box<dyn ViewPartSystem> + Send + Sync>,
+    used_by: HashSet<SpaceViewClassIdentifier>,
+
+    /// Handle to subscription of [`VisualizerEntitySubscriber`] for this visualizer.
+    entity_subscriber_handle: re_arrow_store::StoreSubscriberHandle,
+}
+
+impl Drop for VisualizerTypeRegistryEntry {
+    fn drop(&mut self) {
+        // TODO(andreas): DataStore unsubscribe is not yet implemented!
+        //DataStore::unregister_subscriber(self.entity_subscriber_handle);
+    }
+}
+
 /// Registry of all known space view types.
 ///
 /// Expected to be populated on viewer startup.
 #[derive(Default)]
 pub struct SpaceViewClassRegistry {
     space_view_classes: HashMap<SpaceViewClassIdentifier, SpaceViewClassRegistryEntry>,
-    visualizers: HashMap<ViewSystemIdentifier, SystemTypeRegistryEntry<dyn ViewPartSystem>>,
-    context_systems: HashMap<ViewSystemIdentifier, SystemTypeRegistryEntry<dyn ViewContextSystem>>,
+    context_systems: HashMap<ViewSystemIdentifier, ContextSystemTypeRegistryEntry>,
+    visualizers: HashMap<ViewSystemIdentifier, VisualizerTypeRegistryEntry>,
     placeholder: SpaceViewClassRegistryEntry,
 }
 
@@ -175,8 +204,8 @@ impl SpaceViewClassRegistry {
                 identifier,
                 SpaceViewClassRegistryEntry {
                     class,
-                    context_systems,
-                    visualizers,
+                    context_system_ids: context_systems,
+                    visualizer_system_ids: visualizers,
                 },
             )
             .is_some()
@@ -245,10 +274,25 @@ impl SpaceViewClassRegistry {
     }
 
     /// Iterates over all registered Space View class types.
-    pub fn iter_classes(&self) -> impl Iterator<Item = &dyn DynSpaceViewClass> {
-        self.space_view_classes
-            .values()
-            .map(|entry| entry.class.as_ref())
+    pub fn iter_registry(&self) -> impl Iterator<Item = &SpaceViewClassRegistryEntry> {
+        self.space_view_classes.values()
+    }
+
+    /// Returns the set of entities that are applicable to the given visualizer.
+    ///
+    /// The list is kept up to date by a store subscriber.
+    pub fn applicable_entities_for_visualizer_system(
+        &self,
+        visualizer: ViewSystemIdentifier,
+        store: &re_log_types::StoreId,
+    ) -> Option<IntSet<EntityPath>> {
+        self.visualizers.get(&visualizer).and_then(|entry| {
+            DataStore::with_subscriber::<VisualizerEntitySubscriber, _, _>(
+                entry.entity_subscriber_handle,
+                |subscriber| subscriber.entities(store).cloned(),
+            )
+            .flatten()
+        })
     }
 
     pub fn new_context_collection(
@@ -266,7 +310,7 @@ impl SpaceViewClassRegistry {
 
         ViewContextCollection {
             systems: class
-                .context_systems
+                .context_system_ids
                 .iter()
                 .filter_map(|name| {
                     self.context_systems.get(name).map(|entry| {
@@ -293,7 +337,7 @@ impl SpaceViewClassRegistry {
 
         ViewPartCollection {
             systems: class
-                .visualizers
+                .visualizer_system_ids
                 .iter()
                 .filter_map(|name| {
                     self.visualizers.get(name).map(|entry| {

crates/re_viewer_context/src/space_view/visualizer_entity_subscriber.rs

@@ -0,0 +1,136 @@
+use ahash::HashMap;
+use bit_vec::BitVec;
+use nohash_hasher::{IntMap, IntSet};
+
+use re_arrow_store::StoreSubscriber;
+use re_log_types::{EntityPath, EntityPathHash, StoreId};
+use re_types::ComponentName;
+
+use crate::{IdentifiedViewSystem, ViewPartSystem, ViewSystemIdentifier};
+
+/// A store subscriber that keep track which entities in a store can be
+/// processed by a single given visualizer type.
+///
+/// The list of entities is additive:
+/// If an entity was at any point in time applicable to the visualizer, it will be
+/// kept in the list of entities.
+///
+/// Applicability is determined by the visualizer's set of required components.
+///
+/// There's only a single entity subscriber per visualizer *type*.
+/// This means that if the same visualizer is used in multiple space views, only a single
+/// `VisualizerEntitySubscriber` is created for all of them.
+pub struct VisualizerEntitySubscriber {
+    /// Visualizer type this subscriber is associated with.
+    visualizer: ViewSystemIdentifier,
+
+    /// Assigns each required component an index.
+    required_components_indices: IntMap<ComponentName, usize>,
+
+    per_store_mapping: HashMap<StoreId, VisualizerEntityMapping>,
+}
+
+#[derive(Default)]
+struct VisualizerEntityMapping {
+    /// For each entity, which of the required components are present.
+    ///
+    /// In order of `required_components`.
+    /// If all bits are set, the entity is applicable to the visualizer.
+    // TODO(andreas): We could just limit the number of required components to 32 or 64 and
+    // then use a single u32/u64 as a bitmap.
+    required_component_bitmap_per_entity: IntMap<EntityPathHash, BitVec>,
+
+    /// Which entities the visualizer can be applied to.
+    ///
+    /// Guaranteed to not have any duplicate entries.
+    /// Order is not defined.
+    // TODO(andreas): It would be nice if these were just `EntityPathHash`.
+    applicable_entities: IntSet<EntityPath>,
+}
+
+impl VisualizerEntitySubscriber {
+    pub fn new<T: IdentifiedViewSystem + ViewPartSystem>(visualizer: &T) -> Self {
+        Self {
+            visualizer: T::identifier(),
+            required_components_indices: visualizer
+                .required_components()
+                .into_iter()
+                .enumerate()
+                .map(|(i, name)| (name, i))
+                .collect(),
+            per_store_mapping: Default::default(),
+        }
+    }
+
+    /// List of entities that are applicable to the visualizer.
+    #[inline]
+    pub fn entities(&self, store: &StoreId) -> Option<&IntSet<EntityPath>> {
+        self.per_store_mapping
+            .get(store)
+            .map(|mapping| &mapping.applicable_entities)
+    }
+}
+
+impl StoreSubscriber for VisualizerEntitySubscriber {
+    #[inline]
+    fn name(&self) -> String {
+        self.visualizer.as_str().to_owned()
+    }
+
+    #[inline]
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    #[inline]
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+
+    fn on_events(&mut self, events: &[re_arrow_store::StoreEvent]) {
+        // TODO(andreas): Need to react to store removals as well. As of writing doesn't exist yet.
+
+        for event in events {
+            if event.diff.kind != re_arrow_store::StoreDiffKind::Addition {
+                // Applicability is only additive, don't care about removals.
+                continue;
+            }
+
+            let store_mapping = self
+                .per_store_mapping
+                .entry(event.store_id.clone())
+                .or_default();
+
+            let required_components_bitmap = store_mapping
+                .required_component_bitmap_per_entity
+                .entry(event.diff.entity_path.hash())
+                .or_insert_with(|| {
+                    BitVec::from_elem(self.required_components_indices.len(), false)
+                });
+
+            if required_components_bitmap.all() {
+                // We already know that this entity is applicable to the visualizer.
+                continue;
+            }
+
+            for component_name in event.diff.cells.keys() {
+                if let Some(index) = self.required_components_indices.get(component_name) {
+                    required_components_bitmap.set(*index, true);
+                }
+            }
+
+            if required_components_bitmap.all() {
+                re_log::debug!(
+                    "Entity {:?} in store {:?} is now applicable to visualizer {:?}",
+                    event.diff.entity_path,
+                    event.store_id,
+                    self.visualizer
+                );
+
+                store_mapping
+                    .applicable_entities
+                    .insert(event.diff.entity_path.clone());
+            }
+        }
+    }
+}