diff --git a/Cargo.lock b/Cargo.lock index e186056d4287..4a5e02bb60b3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5287,6 +5287,7 @@ dependencies = [ "ahash 0.8.6", "anyhow", "arboard", + "bit-vec", "bytemuck", "egui", "egui-wgpu", diff --git a/Cargo.toml b/Cargo.toml index 50842815dd5d..3334bc46f2b9 100644 --- a/Cargo.toml +++ b/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" diff --git a/crates/re_viewer_context/Cargo.toml b/crates/re_viewer_context/Cargo.toml index 2d20cf1a4cbe..8101db92153d 100644 --- a/crates/re_viewer_context/Cargo.toml +++ b/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 diff --git a/crates/re_viewer_context/src/space_view/mod.rs b/crates/re_viewer_context/src/space_view/mod.rs index ea00772759e9..2063e4efc492 100644 --- a/crates/re_viewer_context/src/space_view/mod.rs +++ b/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::{ diff --git a/crates/re_viewer_context/src/space_view/space_view_class_registry.rs b/crates/re_viewer_context/src/space_view/space_view_class_registry.rs index d7b3bd4985ba..96f4a41f099c 100644 --- a/crates/re_viewer_context/src/space_view/space_view_class_registry.rs +++ b/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::::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::::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::::default()), + used_by: Default::default(), + entity_subscriber_handle, + } }) .used_by .insert(self.identifier); @@ -110,10 +123,10 @@ impl SpaceViewSystemRegistrator<'_> { } /// Space view class entry in [`SpaceViewClassRegistry`]. -struct SpaceViewClassRegistryEntry { - class: Box, - context_systems: HashSet, - visualizers: HashSet, +pub struct SpaceViewClassRegistryEntry { + pub class: Box, + pub context_system_ids: HashSet, + pub visualizer_system_ids: HashSet, } #[allow(clippy::derivable_impls)] // Clippy gets this one wrong. @@ -121,26 +134,42 @@ impl Default for SpaceViewClassRegistryEntry { fn default() -> Self { Self { class: Box::::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 { - factory_method: Box Box + Send + Sync>, +/// Context system type entry in [`SpaceViewClassRegistry`]. +struct ContextSystemTypeRegistryEntry { + factory_method: Box Box + Send + Sync>, used_by: HashSet, } +/// Visualizer entry in [`SpaceViewClassRegistry`]. +struct VisualizerTypeRegistryEntry { + factory_method: Box Box + Send + Sync>, + used_by: HashSet, + + /// 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, - visualizers: HashMap>, - context_systems: HashMap>, + context_systems: HashMap, + visualizers: HashMap, 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 { - self.space_view_classes - .values() - .map(|entry| entry.class.as_ref()) + pub fn iter_registry(&self) -> impl Iterator { + 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> { + self.visualizers.get(&visualizer).and_then(|entry| { + DataStore::with_subscriber::( + 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| { diff --git a/crates/re_viewer_context/src/space_view/visualizer_entity_subscriber.rs b/crates/re_viewer_context/src/space_view/visualizer_entity_subscriber.rs new file mode 100644 index 000000000000..49b919c0d900 --- /dev/null +++ b/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, + + per_store_mapping: HashMap, +} + +#[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, + + /// 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, +} + +impl VisualizerEntitySubscriber { + pub fn new(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> { + 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()); + } + } + } +} diff --git a/crates/re_viewport/src/space_view_heuristics.rs b/crates/re_viewport/src/space_view_heuristics.rs index 304ae1e572d0..0f426a937096 100644 --- a/crates/re_viewport/src/space_view_heuristics.rs +++ b/crates/re_viewport/src/space_view_heuristics.rs @@ -10,13 +10,11 @@ use re_types::components::{DisconnectedSpace, TensorData}; use re_types::ComponentNameSet; use re_viewer_context::{ AutoSpawnHeuristic, DataQueryResult, EntitiesPerSystem, EntitiesPerSystemPerClass, - HeuristicFilterContext, PerSystemEntities, SpaceViewClassIdentifier, ViewContextCollection, - ViewPartCollection, ViewSystemIdentifier, ViewerContext, + HeuristicFilterContext, PerSystemEntities, SpaceViewClassIdentifier, ViewPartCollection, + ViewerContext, }; -use tinyvec::TinyVec; -use crate::query_pinhole; -use crate::{space_info::SpaceInfoCollection, space_view::SpaceViewBlueprint}; +use crate::{query_pinhole, space_info::SpaceInfoCollection, space_view::SpaceViewBlueprint}; // --------------------------------------------------------------------------- // TODO(#3079): Knowledge of specific space view classes should not leak here. @@ -80,9 +78,9 @@ pub fn all_possible_space_views( // should not influence the heuristics. let entities_used_by_any_part_system_of_class: IntMap<_, _> = ctx .space_view_class_registry - .iter_classes() - .map(|class| { - let class_identifier = class.identifier(); + .iter_registry() + .map(|entry| { + let class_identifier = entry.class.identifier(); let parts = ctx .space_view_class_registry .new_part_collection(class_identifier); @@ -559,107 +557,61 @@ pub fn identify_entities_per_system_per_class( ) -> EntitiesPerSystemPerClass { re_tracing::profile_function!(); - let system_collections_per_class: IntMap< - SpaceViewClassIdentifier, - (ViewContextCollection, ViewPartCollection), - > = space_view_class_registry - .iter_classes() - .map(|class| { - let class_identifier = class.identifier(); - ( - class_identifier, - ( - space_view_class_registry.new_context_collection(class_identifier), - space_view_class_registry.new_part_collection(class_identifier), - ), - ) - }) - .collect(); - - let systems_per_required_components = { - re_tracing::profile_scope!("gather required components per systems"); - - let mut systems_per_required_components: HashMap< - ComponentNameSet, - IntMap>, - > = HashMap::default(); - for (class_identifier, (_context_collection, part_collection)) in - &system_collections_per_class - { - for (system_identifier, part) in part_collection.iter_with_identifiers() { - systems_per_required_components - .entry(part.required_components().into_iter().collect()) - .or_default() - .entry(*class_identifier) - .or_default() - .push(system_identifier); - } - // TODO(#4377): Handle context systems but keep them parallel - /* - for (system_name, part) in context_collection.iter_with_names() { - for components in part.compatible_component_sets() { - systems_per_required_components - .entry(components.into_iter().collect()) - .or_default() - .entry(*class_identifier) - .or_default() - .push(system_name); - } - } - */ - } - systems_per_required_components - }; - - let mut entities_per_system_per_class = EntitiesPerSystemPerClass::default(); - + let store = store_db.store().id(); let heuristic_context = compute_heuristic_context_for_entities(store_db); - let store = store_db.store(); - for ent_path in store_db.entity_paths() { - let Some(components) = store.all_components(&re_log_types::Timeline::log_time(), ent_path) - else { - continue; - }; - - let all_components: ComponentNameSet = components.into_iter().collect(); - - for (required_components, systems_per_class) in &systems_per_required_components { - if !all_components.is_superset(required_components) { - continue; - } - for (class, systems) in systems_per_class { - let Some((_, part_collection)) = system_collections_per_class.get(class) else { - continue; + space_view_class_registry + .iter_registry() + .map(|entry| { + let mut entities_per_system = EntitiesPerSystem::default(); + let class_id = entry.class.identifier(); + + // TODO(andreas): Once `heuristic_filter` is no longer applied, we don't need to instantiate the systems anymore. + //for system_id in &entry.visualizer_system_ids { + for (system_id, system) in space_view_class_registry + .new_part_collection(class_id) + .systems + { + let entities: IntSet = if let Some(entities) = space_view_class_registry + .applicable_entities_for_visualizer_system(system_id, store) + { + // TODO(andreas): Don't apply heuristic_filter here, this should be part of the query! + // Note that once this is done, `EntitiesPerSystemPerClass` becomes just `EntitiesPerSystem` since there's never a need to distinguish per class! + entities + .into_iter() + .filter(|ent_path| { + let Some(components) = store_db + .store() + .all_components(&re_log_types::Timeline::log_time(), ent_path) + else { + return false; + }; + + let all_components: ComponentNameSet = components.into_iter().collect(); + + system.heuristic_filter( + store_db.store(), + ent_path, + heuristic_context + .get(ent_path) + .copied() + .unwrap_or_default() + .with_class(class_id), + current_query, + &all_components, + ) + }) + .collect() + } else { + Default::default() }; - for system in systems { - if let Ok(view_part_system) = part_collection.get_by_identifier(*system) { - if !view_part_system.heuristic_filter( - store, - ent_path, - heuristic_context - .get(ent_path) - .copied() - .unwrap_or_default() - .with_class(*class), - current_query, - &all_components, - ) { - continue; - } - } - - entities_per_system_per_class - .entry(*class) - .or_default() - .entry(*system) - .or_default() - .insert(ent_path.clone()); - } + entities_per_system.insert(system_id, entities); } - } - } - entities_per_system_per_class + // TODO(#4377): Handle context systems but keep them parallel + + (class_id, entities_per_system) + }) + .collect() } diff --git a/crates/re_viewport/src/viewport_blueprint_ui.rs b/crates/re_viewport/src/viewport_blueprint_ui.rs index fe381759438f..44dffa4ab821 100644 --- a/crates/re_viewport/src/viewport_blueprint_ui.rs +++ b/crates/re_viewport/src/viewport_blueprint_ui.rs @@ -434,14 +434,14 @@ impl Viewport<'_, '_> { // Empty space views of every available types for space_view in ctx .space_view_class_registry - .iter_classes() - .sorted_by_key(|space_view_class| space_view_class.display_name()) - .map(|class| { + .iter_registry() + .sorted_by_key(|entry| entry.class.display_name()) + .map(|entry| { SpaceViewBlueprint::new( - class.identifier(), - &format!("empty {}", class.display_name()), + entry.class.identifier(), + &format!("empty {}", entry.class.display_name()), &EntityPath::root(), - DataQueryBlueprint::new(class.identifier(), std::iter::empty()), + DataQueryBlueprint::new(entry.class.identifier(), std::iter::empty()), ) }) {