Dynamically registered space view (part/context) systems

crates/re_data_ui/src/lib.rs

@@ -156,15 +156,6 @@ pub fn annotations(
 ) -> std::sync::Arc<re_viewer_context::Annotations> {
     re_tracing::profile_function!();
     let mut annotation_map = re_viewer_context::AnnotationMap::default();
-    let entity_paths: nohash_hasher::IntSet<_> = std::iter::once(entity_path.clone()).collect();
-    let entity_props_map = re_data_store::EntityPropertyMap::default();
-    let query = re_viewer_context::ViewQuery::new(
-        entity_path,
-        &entity_paths,
-        query.timeline,
-        query.at,
-        &entity_props_map,
-    );
-    annotation_map.load(ctx, &query);
+    annotation_map.load(ctx, query, std::iter::once(entity_path));
     annotation_map.find(entity_path)
 }

crates/re_space_view_bar_chart/src/space_view_class.rs

@@ -3,7 +3,9 @@ use re_components::TensorData;
 use re_log_types::EntityPath;
 use re_space_view::controls;
 use re_viewer_context::{
-    auto_color, SpaceViewClass, SpaceViewClassName, SpaceViewId, TypedScene, ViewerContext,
+    auto_color, SpaceViewClass, SpaceViewClassName, SpaceViewClassRegistryError, SpaceViewId,
+    SpaceViewSystemExecutionError, ViewContextCollection, ViewPartCollection, ViewQuery,
+    ViewerContext,
 };
 
 use super::view_part_system::BarChartViewPartSystem;
@@ -13,9 +15,6 @@ pub struct BarChartSpaceView;
 
 impl SpaceViewClass for BarChartSpaceView {
     type State = ();
-    type Context = ();
-    type SystemCollection = BarChartViewPartSystem;
-    type ViewPartData = ();
 
     fn name(&self) -> SpaceViewClassName {
         "Bar Chart".into()
@@ -46,6 +45,13 @@ impl SpaceViewClass for BarChartSpaceView {
         layout.layout_job.into()
     }
 
+    fn on_register(
+        &self,
+        system_registry: &mut re_viewer_context::SpaceViewSystemRegistry,
+    ) -> Result<(), SpaceViewClassRegistryError> {
+        system_registry.register_part_system::<BarChartViewPartSystem>()
+    }
+
     fn preferred_tile_aspect_ratio(&self, _state: &Self::State) -> Option<f32> {
         None
     }
@@ -69,12 +75,15 @@ impl SpaceViewClass for BarChartSpaceView {
         _ctx: &mut ViewerContext<'_>,
         ui: &mut egui::Ui,
         _state: &mut Self::State,
-        scene: &mut TypedScene<Self>,
-        _space_origin: &EntityPath,
-        _space_view_id: SpaceViewId,
-    ) {
+        _view_ctx: &ViewContextCollection,
+        parts: &ViewPartCollection,
+        _query: &ViewQuery<'_>,
+        _draw_data: Vec<re_renderer::QueueableDrawData>,
+    ) -> Result<(), SpaceViewSystemExecutionError> {
         use egui::plot::{Bar, BarChart, Legend, Plot};
 
+        let charts = &parts.get::<BarChartViewPartSystem>()?.charts;
+
         ui.scope(|ui| {
             Plot::new("bar_chart_plot")
                 .legend(Legend::default())
@@ -102,7 +111,7 @@ impl SpaceViewClass for BarChartSpaceView {
                         .color(color)
                     }
 
-                    for (ent_path, tensor) in &scene.parts.charts {
+                    for (ent_path, tensor) in charts.iter() {
                         let chart = match &tensor.data {
                             TensorData::U8(data) => {
                                 create_bar_chart(ent_path, data.iter().copied())
@@ -150,5 +159,7 @@ impl SpaceViewClass for BarChartSpaceView {
                     }
                 });
         });
+
+        Ok(())
     }
 }

crates/re_space_view_bar_chart/src/view_part_system.rs

@@ -5,31 +5,27 @@ use re_components::Tensor;
 use re_data_store::EntityPath;
 use re_log_types::Component as _;
 use re_viewer_context::{
-    ArchetypeDefinition, SpaceViewClass, SpaceViewHighlights, ViewPartSystem, ViewQuery,
-    ViewerContext,
+    ArchetypeDefinition, SpaceViewSystemExecutionError, ViewContextCollection, ViewPartSystem,
+    ViewQuery, ViewerContext,
 };
 
-use crate::BarChartSpaceView;
-
 /// A bar chart system, with everything needed to render it.
 #[derive(Default)]
 pub struct BarChartViewPartSystem {
     pub charts: BTreeMap<EntityPath, Tensor>,
 }
 
-impl ViewPartSystem<BarChartSpaceView> for BarChartViewPartSystem {
+impl ViewPartSystem for BarChartViewPartSystem {
     fn archetype(&self) -> ArchetypeDefinition {
         vec1::vec1![Tensor::name()]
     }
 
-    fn populate(
+    fn execute(
         &mut self,
         ctx: &mut ViewerContext<'_>,
         query: &ViewQuery<'_>,
-        _state: &<BarChartSpaceView as SpaceViewClass>::State,
-        _context: &<BarChartSpaceView as SpaceViewClass>::Context,
-        _highlights: &SpaceViewHighlights,
-    ) -> Vec<re_renderer::QueueableDrawData> {
+        _view_ctx: &ViewContextCollection,
+    ) -> Result<Vec<re_renderer::QueueableDrawData>, SpaceViewSystemExecutionError> {
         re_tracing::profile_function!();
 
         let store = &ctx.store_db.entity_db.data_store;
@@ -49,6 +45,10 @@ impl ViewPartSystem<BarChartSpaceView> for BarChartViewPartSystem {
             }
         }
 
-        Vec::new()
+        Ok(Vec::new())
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
     }
 }

crates/re_space_view_spatial/src/contexts/annotation_context.rs

@@ -10,12 +10,19 @@ impl ViewContextSystem for AnnotationSceneContext {
         vec![vec1::vec1![AnnotationContext::name()]]
     }
 
-    fn populate(
+    fn execute(
         &mut self,
         ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
-        _space_view_state: &dyn re_viewer_context::SpaceViewState,
     ) {
-        self.0.load(ctx, query);
+        self.0.load(
+            ctx,
+            &query.latest_at_query(),
+            query.iter_entities().map(|(p, _)| p),
+        );
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
     }
 }

crates/re_space_view_spatial/src/contexts/depth_offsets.rs

@@ -2,7 +2,7 @@ use std::collections::{BTreeMap, BTreeSet};
 
 use nohash_hasher::IntMap;
 use re_components::DrawOrder;
-use re_log_types::{Component, EntityPath, EntityPathHash};
+use re_log_types::{Component, EntityPathHash};
 use re_viewer_context::{ArchetypeDefinition, ViewContextSystem};
 
 /// Context for creating a mapping from [`DrawOrder`] to [`re_renderer::DepthOffset`].
@@ -23,11 +23,10 @@ impl ViewContextSystem for EntityDepthOffsets {
         vec![vec1::vec1![DrawOrder::name()]]
     }
 
-    fn populate(
+    fn execute(
         &mut self,
         ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
-        _space_view_state: &dyn re_viewer_context::SpaceViewState,
     ) {
         re_tracing::profile_function!();
 
@@ -117,10 +116,8 @@ impl ViewContextSystem for EntityDepthOffsets {
             })
             .collect();
     }
-}
 
-impl EntityDepthOffsets {
-    pub fn get(&self, ent_path: &EntityPath) -> Option<re_renderer::DepthOffset> {
-        self.per_entity.get(&ent_path.hash()).cloned()
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
     }
 }

crates/re_space_view_spatial/src/contexts/mod.rs

@@ -8,84 +8,58 @@ use std::sync::atomic::AtomicUsize;
 
 pub use annotation_context::AnnotationSceneContext;
 pub use depth_offsets::EntityDepthOffsets;
+pub use non_interactive_entities::NonInteractiveEntities;
 pub use shared_render_builders::SharedRenderBuilders;
 pub use transform_context::{pinhole_camera_view_coordinates, TransformContext};
 
 // -----------------------------------------------------------------------------
 
-use re_log_types::EntityPath;
 use re_renderer::DepthOffset;
-use re_viewer_context::{Annotations, ViewContext};
+use re_viewer_context::{Annotations, SpaceViewClassRegistryError, ViewContextSystem};
 
-use self::non_interactive_entities::NonInteractiveEntities;
+/// Context objects for a single entity in a spatial scene.
+pub struct SpatialSceneEntityContext<'a> {
+    pub world_from_obj: glam::Affine3A,
+    pub depth_offset: DepthOffset,
+    pub annotations: std::sync::Arc<Annotations>,
+    pub shared_render_builders: &'a SharedRenderBuilders,
+    pub counter: &'a PrimitiveCounter,
 
-#[derive(Default)]
-pub struct SpatialViewContext {
-    pub transforms: TransformContext,
-    pub depth_offsets: EntityDepthOffsets,
-    pub annotations: AnnotationSceneContext,
-    pub shared_render_builders: SharedRenderBuilders,
-    pub non_interactive_entities: NonInteractiveEntities,
+    pub highlight: &'a re_viewer_context::SpaceViewOutlineMasks, // Not part of the context, but convenient to have here.
+}
 
+// TODO(andreas): num_3d_primitives is likely not needed in this form once 2D/3D separation is gone. num_primitives should be handled in a more general way?
+#[derive(Default)]
+pub struct PrimitiveCounter {
     pub num_primitives: AtomicUsize,
     pub num_3d_primitives: AtomicUsize,
 }
 
-impl ViewContext for SpatialViewContext {
-    fn vec_mut(&mut self) -> Vec<&mut dyn re_viewer_context::ViewContextSystem> {
-        let Self {
-            transforms,
-            depth_offsets,
-            annotations,
-            shared_render_builders,
-            non_interactive_entities,
-            num_3d_primitives: _,
-            num_primitives: _,
-        } = self;
-        vec![
-            transforms,
-            depth_offsets,
-            annotations,
-            shared_render_builders,
-            non_interactive_entities,
-        ]
+impl ViewContextSystem for PrimitiveCounter {
+    fn archetypes(&self) -> Vec<re_viewer_context::ArchetypeDefinition> {
+        Vec::new()
     }
 
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
+    fn execute(
+        &mut self,
+        _ctx: &mut re_viewer_context::ViewerContext<'_>,
+        _query: &re_viewer_context::ViewQuery<'_>,
+    ) {
     }
 
-    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+    fn as_any(&self) -> &dyn std::any::Any {
         self
     }
 }
 
-impl SpatialViewContext {
-    pub fn lookup_entity_context<'a>(
-        &'a self,
-        ent_path: &EntityPath,
-        highlights: &'a re_viewer_context::SpaceViewHighlights,
-        default_depth_offset: DepthOffset,
-    ) -> Option<SpatialSceneEntityContext<'a>> {
-        Some(SpatialSceneEntityContext {
-            world_from_obj: self.transforms.reference_from_entity(ent_path)?,
-            depth_offset: *self
-                .depth_offsets
-                .per_entity
-                .get(&ent_path.hash())
-                .unwrap_or(&default_depth_offset),
-            annotations: self.annotations.0.find(ent_path),
-            shared_render_builders: &self.shared_render_builders,
-            highlight: highlights.entity_outline_mask(ent_path.hash()),
-        })
-    }
-}
-
-/// Context objects for a single entity in a spatial scene.
-pub struct SpatialSceneEntityContext<'a> {
-    pub world_from_obj: glam::Affine3A,
-    pub depth_offset: DepthOffset,
-    pub annotations: std::sync::Arc<Annotations>,
-    pub shared_render_builders: &'a SharedRenderBuilders,
-    pub highlight: &'a re_viewer_context::SpaceViewOutlineMasks, // Not part of the context, but convenient to have here.
+pub fn register_contexts(
+    system_registry: &mut re_viewer_context::SpaceViewSystemRegistry,
+) -> Result<(), SpaceViewClassRegistryError> {
+    system_registry.register_context_system::<TransformContext>()?;
+    system_registry.register_context_system::<EntityDepthOffsets>()?;
+    system_registry.register_context_system::<AnnotationSceneContext>()?;
+    system_registry.register_context_system::<SharedRenderBuilders>()?;
+    system_registry.register_context_system::<NonInteractiveEntities>()?;
+    system_registry.register_context_system::<PrimitiveCounter>()?;
+    Ok(())
 }

crates/re_space_view_spatial/src/contexts/non_interactive_entities.rs

@@ -13,11 +13,10 @@ impl ViewContextSystem for NonInteractiveEntities {
         Vec::new()
     }
 
-    fn populate(
+    fn execute(
         &mut self,
         _ctx: &mut re_viewer_context::ViewerContext<'_>,
         query: &re_viewer_context::ViewQuery<'_>,
-        _space_view_state: &dyn re_viewer_context::SpaceViewState,
     ) {
         re_tracing::profile_function!();
 
@@ -33,4 +32,8 @@ impl ViewContextSystem for NonInteractiveEntities {
             })
             .collect();
     }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
 }