Extend/iterate on SpaceViewClass framework with SceneContext & port SpatialSpaceView scene parts

crates/re_data_ui/src/lib.rs

@@ -159,6 +159,7 @@ pub fn annotations(
     let entity_paths: nohash_hasher::IntSet<_> = std::iter::once(entity_path.clone()).collect();
     let entity_props_map = re_data_store::EntityPropertyMap::default();
     let scene_query = re_viewer_context::SceneQuery {
+        space_path: entity_path,
         entity_paths: &entity_paths,
         timeline: query.timeline,
         latest_at: query.at,

crates/re_space_view/src/empty_scene_context.rs

@@ -0,0 +1,19 @@
+use re_viewer_context::{SceneContext, SceneContextPart};
+
+/// Implementation of an empty scene context.
+#[derive(Default)]
+pub struct EmptySceneContext;
+
+impl SceneContext for EmptySceneContext {
+    fn vec_mut(&mut self) -> Vec<&mut dyn SceneContextPart> {
+        Vec::new()
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+}

crates/re_space_view/src/empty_space_view_state.rs

@@ -0,0 +1,15 @@
+use re_viewer_context::SpaceViewState;
+
+/// Space view state without any contents.
+#[derive(Default)]
+pub struct EmptySpaceViewState;
+
+impl SpaceViewState for EmptySpaceViewState {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+}

crates/re_space_view/src/lib.rs

@@ -4,9 +4,11 @@
 
 pub mod controls;
 mod data_blueprint;
-mod highlights;
+mod empty_scene_context;
+mod empty_space_view_state;
 mod screenshot;
 
 pub use data_blueprint::{DataBlueprintGroup, DataBlueprintTree};
-pub use highlights::{SpaceViewEntityHighlight, SpaceViewHighlights, SpaceViewOutlineMasks};
+pub use empty_scene_context::EmptySceneContext;
+pub use empty_space_view_state::EmptySpaceViewState;
 pub use screenshot::ScreenshotMode;

crates/re_space_view_spatial/Cargo.toml

@@ -34,12 +34,14 @@ re_space_view.workspace = true
 ahash.workspace = true
 anyhow.workspace = true
 bytemuck.workspace = true
-egui.workspace = true
 eframe.workspace = true
+egui.workspace = true
 glam.workspace = true
 itertools.workspace = true
 macaw = { workspace = true, features = ["with_serde"] }
 nohash-hasher.workspace = true
+parking_lot.workspace = true
 serde = "1"
 smallvec = { workspace = true, features = ["serde"] }
+vec1.workspace = true
 wgpu.workspace = true

crates/re_space_view_spatial/src/lib.rs

@@ -8,14 +8,12 @@ mod mesh_cache;
 mod mesh_loader;
 mod scene;
 mod space_camera_3d;
-mod transform_cache;
 mod ui;
 mod ui_2d;
 mod ui_3d;
 mod ui_renderer_bridge;
 
 // TODO(andreas) should only make the main type public
 
-pub use scene::SceneSpatial;
-pub use transform_cache::{TransformCache, UnreachableTransform};
+pub use scene::{SceneSpatial, TransformContext, UnreachableTransform};
 pub use ui::{SpatialNavigationMode, ViewSpatialState};

crates/re_space_view_spatial/src/scene/contexts/annotation_context.rs

@@ -0,0 +1,21 @@
+use re_components::AnnotationContext;
+use re_log_types::Component;
+use re_viewer_context::{AnnotationMap, ArchetypeDefinition, SceneContextPart};
+
+#[derive(Default)]
+pub struct AnnotationSceneContext(pub AnnotationMap);
+
+impl SceneContextPart for AnnotationSceneContext {
+    fn archetypes(&self) -> Vec<ArchetypeDefinition> {
+        vec![vec1::vec1![AnnotationContext::name()]]
+    }
+
+    fn populate(
+        &mut self,
+        ctx: &mut re_viewer_context::ViewerContext<'_>,
+        query: &re_viewer_context::SceneQuery<'_>,
+        _space_view_state: &dyn re_viewer_context::SpaceViewState,
+    ) {
+        self.0.load(ctx, query);
+    }
+}

crates/re_space_view_spatial/src/scene/contexts/depth_offsets.rs

@@ -0,0 +1,118 @@
+use std::collections::{BTreeMap, BTreeSet};
+
+use nohash_hasher::IntMap;
+use re_components::DrawOrder;
+use re_log_types::{Component, EntityPathHash};
+use re_viewer_context::{ArchetypeDefinition, SceneContextPart};
+
+// TODO(wumpf): In the new traitified system it would be likely more ergonomic to map from darw order to depth offset
+//              This would also yield more consistent/meaningful archetypes and avoid queries where not needed.
+#[derive(Default)]
+pub struct EntityDepthOffsets {
+    pub per_entity: IntMap<EntityPathHash, re_renderer::DepthOffset>,
+    pub box2d: re_renderer::DepthOffset,
+    pub lines2d: re_renderer::DepthOffset,
+    pub image: re_renderer::DepthOffset,
+    pub points: re_renderer::DepthOffset,
+}
+
+impl SceneContextPart for EntityDepthOffsets {
+    fn archetypes(&self) -> Vec<ArchetypeDefinition> {
+        vec![vec1::vec1![DrawOrder::name()]]
+    }
+
+    fn populate(
+        &mut self,
+        ctx: &mut re_viewer_context::ViewerContext<'_>,
+        query: &re_viewer_context::SceneQuery<'_>,
+        _space_view_state: &dyn re_viewer_context::SpaceViewState,
+    ) {
+        re_tracing::profile_function!();
+
+        #[derive(PartialEq, PartialOrd, Eq, Ord)]
+        enum DrawOrderTarget {
+            Entity(EntityPathHash),
+            DefaultBox2D,
+            DefaultLines2D,
+            DefaultImage,
+            DefaultPoints,
+        }
+
+        let store = &ctx.store_db.entity_db.data_store;
+
+        // 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() {
+            if let Some(draw_order) = store.query_latest_component::<DrawOrder>(
+                ent_path,
+                &ctx.rec_cfg.time_ctrl.current_query(),
+            ) {
+                entities_per_draw_order
+                    .entry(draw_order)
+                    .or_default()
+                    .insert(DrawOrderTarget::Entity(ent_path.hash()));
+            }
+        }
+
+        // Push in default draw orders. All of them using the none hash.
+        entities_per_draw_order.insert(
+            DrawOrder::DEFAULT_BOX2D,
+            [DrawOrderTarget::DefaultBox2D].into(),
+        );
+        entities_per_draw_order.insert(
+            DrawOrder::DEFAULT_IMAGE,
+            [DrawOrderTarget::DefaultImage].into(),
+        );
+        entities_per_draw_order.insert(
+            DrawOrder::DEFAULT_LINES2D,
+            [DrawOrderTarget::DefaultLines2D].into(),
+        );
+        entities_per_draw_order.insert(
+            DrawOrder::DEFAULT_POINTS2D,
+            [DrawOrderTarget::DefaultPoints].into(),
+        );
+
+        // Determine re_renderer draw order from this.
+        //
+        // We give objects with the same `DrawOrder` still a different depth offset
+        // in order to avoid z-fighting artifacts when rendering in 3D.
+        // (for pure 2D this isn't necessary)
+        //
+        // We want to be as tightly around 0 as possible.
+        let num_entities_with_draw_order: usize = entities_per_draw_order
+            .values()
+            .map(|entities| entities.len())
+            .sum();
+        let mut draw_order = -((num_entities_with_draw_order / 2) as re_renderer::DepthOffset);
+        self.per_entity = entities_per_draw_order
+            .into_values()
+            .flat_map(|targets| {
+                targets
+                    .into_iter()
+                    .filter_map(|target| {
+                        draw_order += 1;
+                        match target {
+                            DrawOrderTarget::Entity(entity) => Some((entity, draw_order)),
+                            DrawOrderTarget::DefaultBox2D => {
+                                self.box2d = draw_order;
+                                None
+                            }
+                            DrawOrderTarget::DefaultLines2D => {
+                                self.lines2d = draw_order;
+                                None
+                            }
+                            DrawOrderTarget::DefaultImage => {
+                                self.image = draw_order;
+                                None
+                            }
+                            DrawOrderTarget::DefaultPoints => {
+                                self.points = draw_order;
+                                None
+                            }
+                        }
+                    })
+                    .collect::<Vec<_>>()
+            })
+            .collect();
+    }
+}

crates/re_space_view_spatial/src/scene/contexts/mod.rs

@@ -0,0 +1,84 @@
+mod annotation_context;
+mod depth_offsets;
+mod shared_render_builders;
+mod transform_context;
+
+use std::sync::atomic::AtomicUsize;
+
+pub use annotation_context::AnnotationSceneContext;
+pub use depth_offsets::EntityDepthOffsets;
+pub use shared_render_builders::SharedRenderBuilders;
+pub use transform_context::{TransformContext, UnreachableTransform};
+
+// -----------------------------------------------------------------------------
+
+use re_log_types::EntityPath;
+use re_renderer::DepthOffset;
+use re_viewer_context::{Annotations, SceneContext};
+
+#[derive(Default)]
+pub struct SpatialSceneContext {
+    pub transforms: TransformContext,
+    pub depth_offsets: EntityDepthOffsets,
+    pub annotations: AnnotationSceneContext,
+    pub shared_render_builders: SharedRenderBuilders,
+
+    pub num_primitives: AtomicUsize,
+}
+
+impl SceneContext for SpatialSceneContext {
+    fn vec_mut(&mut self) -> Vec<&mut dyn re_viewer_context::SceneContextPart> {
+        let Self {
+            transforms,
+            depth_offsets,
+            annotations,
+            shared_render_builders,
+
+            num_primitives: _,
+        } = self;
+        vec![
+            transforms,
+            depth_offsets,
+            annotations,
+            shared_render_builders,
+        ]
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+}
+
+impl SpatialSceneContext {
+    pub fn query<'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.
+}