Allow for custom visualizer with custom shader that integrates into existing view

Cargo.lock

@@ -2084,6 +2084,15 @@ dependencies = [
  "rerun",
 ]
 
+[[package]]
+name = "custom_visualizer"
+version = "0.24.0-alpha.3+dev"
+dependencies = [
+ "bytemuck",
+ "mimalloc",
+ "rerun",
+]
+
 [[package]]
 name = "darling"
 version = "0.20.10"

crates/viewer/re_view_spatial/src/lib.rs

@@ -32,6 +32,9 @@ pub use ui::SpatialViewState;
 pub use view_2d::SpatialView2D;
 pub use view_3d::SpatialView3D;
 
+// Export some other types that are useful for extensions.
+pub use contexts::TransformTreeContext;
+
 pub(crate) use pickable_textured_rect::{PickableRectSourceData, PickableTexturedRect};
 pub(crate) use pinhole::Pinhole;
 

crates/viewer/re_viewer/src/lib.rs

@@ -42,8 +42,7 @@ pub use re_viewer_context::{
 };
 
 pub mod external {
-    pub use parking_lot;
-    pub use {eframe, egui};
+    pub use {eframe, egui, parking_lot};
     pub use {
         re_chunk, re_chunk::external::*, re_chunk_store, re_chunk_store::external::*, re_data_ui,
         re_entity_db, re_log, re_log_types, re_memory, re_renderer, re_smart_channel, re_types,

crates/viewer/re_viewer_context/src/view/view_class_registry.rs

@@ -31,12 +31,15 @@ pub enum ViewClassRegistryError {
     UnknownClassIdentifier(ViewClassIdentifier),
 }
 
-/// Utility for registering view systems, passed on to [`crate::ViewClass::on_register`].
+/// Utility for registering view systems to a given class, passed on to [`crate::ViewClass::on_register`].
 pub struct ViewSystemRegistrator<'a> {
-    registry: &'a mut ViewClassRegistry,
     identifier: ViewClassIdentifier,
-    context_systems: HashSet<ViewSystemIdentifier>,
-    visualizers: HashSet<ViewSystemIdentifier>,
+
+    all_context_systems: &'a mut HashMap<ViewSystemIdentifier, ContextSystemTypeRegistryEntry>,
+    all_visualizers: &'a mut HashMap<ViewSystemIdentifier, VisualizerTypeRegistryEntry>,
+
+    class_context_system: &'a mut HashSet<ViewSystemIdentifier>,
+    class_visualizer_system: &'a mut HashSet<ViewSystemIdentifier>,
 }
 
 impl ViewSystemRegistrator<'_> {
@@ -50,15 +53,14 @@ impl ViewSystemRegistrator<'_> {
         &mut self,
     ) -> Result<(), ViewClassRegistryError> {
         // Name should not overlap with context systems.
-        if self.registry.visualizers.contains_key(&T::identifier()) {
+        if self.all_visualizers.contains_key(&T::identifier()) {
             return Err(ViewClassRegistryError::IdentifierAlreadyInUseForVisualizer(
                 T::identifier().as_str(),
             ));
         }
 
-        if self.context_systems.insert(T::identifier()) {
-            self.registry
-                .context_systems
+        if self.class_context_system.insert(T::identifier()) {
+            self.all_context_systems
                 .entry(T::identifier())
                 .or_insert_with(|| ContextSystemTypeRegistryEntry {
                     factory_method: Box::new(|| Box::<T>::default()),
@@ -85,17 +87,16 @@ impl ViewSystemRegistrator<'_> {
         &mut self,
     ) -> Result<(), ViewClassRegistryError> {
         // Name should not overlap with context systems.
-        if self.registry.context_systems.contains_key(&T::identifier()) {
+        if self.all_context_systems.contains_key(&T::identifier()) {
             return Err(
                 ViewClassRegistryError::IdentifierAlreadyInUseForContextSystem(
                     T::identifier().as_str(),
                 ),
             );
         }
 
-        if self.visualizers.insert(T::identifier()) {
-            self.registry
-                .visualizers
+        if self.class_visualizer_system.insert(T::identifier()) {
+            self.all_visualizers
                 .entry(T::identifier())
                 .or_insert_with(|| {
                     let entity_subscriber_handle = ChunkStore::register_subscriber(Box::new(
@@ -180,41 +181,84 @@ impl ViewClassRegistry {
     pub fn add_class<T: ViewClass + Default + 'static>(
         &mut self,
     ) -> Result<(), ViewClassRegistryError> {
-        let class = Box::<T>::default();
+        let mut class_entry = ViewClassRegistryEntry {
+            class: Box::<T>::default(),
+            identifier: T::identifier(),
+            context_system_ids: Default::default(),
+            visualizer_system_ids: Default::default(),
+        };
 
         let mut registrator = ViewSystemRegistrator {
-            registry: self,
             identifier: T::identifier(),
-            context_systems: Default::default(),
-            visualizers: Default::default(),
+            all_visualizers: &mut self.visualizers,
+            all_context_systems: &mut self.context_systems,
+            class_context_system: &mut class_entry.context_system_ids,
+            class_visualizer_system: &mut class_entry.visualizer_system_ids,
         };
 
-        class.on_register(&mut registrator)?;
-
-        let ViewSystemRegistrator {
-            registry: _,
-            identifier,
-            context_systems,
-            visualizers,
-        } = registrator;
+        class_entry.class.on_register(&mut registrator)?;
 
         if self
             .view_classes
-            .insert(
-                identifier,
-                ViewClassRegistryEntry {
-                    class,
-                    identifier,
-                    context_system_ids: context_systems,
-                    visualizer_system_ids: visualizers,
-                },
-            )
+            .insert(T::identifier(), class_entry)
             .is_some()
         {
-            return Err(ViewClassRegistryError::DuplicateClassIdentifier(identifier));
+            Err(ViewClassRegistryError::DuplicateClassIdentifier(
+                T::identifier(),
+            ))
+        } else {
+            Ok(())
         }
+    }
 
-        Ok(())
+    /// Registers an additional [`VisualizerSystem`] for a view class.
+    ///
+    /// Usually, visualizers are registered in [`ViewClass::on_register`] which is called when
+    /// the class is first registered.
+    /// This method allows to extend an existing class with new visualizers.
+    pub fn register_visualizer<T: VisualizerSystem + IdentifiedViewSystem + Default + 'static>(
+        &mut self,
+        view_class: ViewClassIdentifier,
+    ) -> Result<(), ViewClassRegistryError> {
+        let Some(class_entry) = self.view_classes.get_mut(&view_class) else {
+            return Err(ViewClassRegistryError::UnknownClassIdentifier(view_class));
+        };
+
+        let mut registrator = ViewSystemRegistrator {
+            identifier: class_entry.identifier,
+            all_visualizers: &mut self.visualizers,
+            all_context_systems: &mut self.context_systems,
+            class_context_system: &mut class_entry.context_system_ids,
+            class_visualizer_system: &mut class_entry.visualizer_system_ids,
+        };
+
+        registrator.register_visualizer::<T>()
+    }
+
+    /// Registers an additional [`ViewContextSystem`] for a view class.
+    ///
+    /// Usually, context systems are registered in [`ViewClass::on_register`] which is called when
+    /// the class is first registered.
+    /// This method allows to extend an existing class with new context systems.
+    pub fn register_context_system<
+        T: ViewContextSystem + IdentifiedViewSystem + Default + 'static,
+    >(
+        &mut self,
+        view_class: ViewClassIdentifier,
+    ) -> Result<(), ViewClassRegistryError> {
+        let Some(class_entry) = self.view_classes.get_mut(&view_class) else {
+            return Err(ViewClassRegistryError::UnknownClassIdentifier(view_class));
+        };
+
+        let mut registrator = ViewSystemRegistrator {
+            identifier: class_entry.identifier,
+            all_visualizers: &mut self.visualizers,
+            all_context_systems: &mut self.context_systems,
+            class_context_system: &mut class_entry.context_system_ids,
+            class_visualizer_system: &mut class_entry.visualizer_system_ids,
+        };
+
+        registrator.register_context_system::<T>()
     }
 
     /// Removes a view class from the registry.

examples/rust/custom_visualizer/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "custom_visualizer"
+version = "0.24.0-alpha.3+dev"
+edition = "2021"
+rust-version = "1.85"
+license = "MIT OR Apache-2.0"
+publish = false
+
+[features]
+default = []
+
+# Turn on if you want to share analytics with Rerun (e.g. callstacks on crashes).
+analytics = ["rerun/analytics"]
+
+[dependencies]
+rerun = { path = "../../../crates/top/rerun", default-features = false, features = [
+  "native_viewer",
+  "run",
+] }
+
+# Can't rely on export from `rerun` crate since proc macros can't be re-exported.
+bytemuck = "1"
+
+# mimalloc is a much faster allocator:
+mimalloc = "0.1.43"

examples/rust/custom_visualizer/README.md

@@ -0,0 +1,32 @@
+TODO: copied from custom_view, need to update everything in here.
+
+<!--[metadata]
+title = "Custom view UI"
+thumbnail = "https://static.rerun.io/custom_space_view/e05a073d64003645b6af6de91b068c2f646c1b8a/480w.jpeg"
+thumbnail_dimensions = [480, 343]
+-->
+
+
+<picture>
+  <source media="(max-width: 480px)" srcset="https://static.rerun.io/api_demo/e05a073d64003645b6af6de91b068c2f646c1b8a/480w.jpeg">
+  <source media="(max-width: 768px)" srcset="https://static.rerun.io/api_demo/e05a073d64003645b6af6de91b068c2f646c1b8a/768w.jpeg">
+  <source media="(max-width: 1024px)" srcset="https://static.rerun.io/api_demo/e05a073d64003645b6af6de91b068c2f646c1b8a/1024w.jpeg">
+  <source media="(max-width: 1200px)" srcset="https://static.rerun.io/api_demo/e05a073d64003645b6af6de91b068c2f646c1b8a/1200w.jpeg">
+  <img src="https://static.rerun.io/api_demo/e05a073d64003645b6af6de91b068c2f646c1b8a/full.jpeg" alt="Custom View UI example screenshot">
+</picture>
+
+Example showing how to add custom View classes to extend the Rerun Viewer.
+
+The example is really basic, but should be something you can build upon.
+
+The example starts an SDK server which the Python or Rust logging SDK can connect to.
+
+
+[#2337](https://github.com/rerun-io/rerun/issues/2337): Note that in order to spawn a web viewer with these customizations applied,
+you have to build the web viewer of the version yourself.
+This is currently not supported outside of the Rerun repository.
+
+## Testing it
+Start it with `cargo run -p custom_view`.
+
+Then put some data into it with: `cargo run -p minimal_options -- --connect`

examples/rust/custom_visualizer/shader/custom.wgsl

@@ -0,0 +1,61 @@
+#import <types.wgsl>
+#import <global_bindings.wgsl>
+
+struct UniformBuffer {
+    world_from_obj: mat4x4f,
+
+    color: vec4f,
+
+    picking_layer_object_id: vec2u,
+    picking_instance_id: vec2u,
+
+    outline_mask: vec2u,
+};
+@group(1) @binding(0)
+var<uniform> ubo: UniformBuffer;
+
+struct VertexOut {
+    @location(0) color: vec4f,
+    @builtin(position) position: vec4f,
+};
+
+var<private> v_positions: array<vec2f, 3> = array<vec2f, 3>(
+    vec2f(0.0, 1.0),
+    vec2f(1.0, -1.0),
+    vec2f(-1.0, -1.0),
+);
+
+var<private> v_colors: array<vec4f, 3> = array<vec4f, 3>(
+    vec4f(1.0, 0.0, 0.0, 1.0),
+    vec4f(0.0, 1.0, 0.0, 1.0),
+    vec4f(0.0, 0.0, 1.0, 1.0),
+);
+
+@vertex
+fn vs_main(@builtin(vertex_index) v_idx: u32) -> VertexOut {
+    var out: VertexOut;
+
+    let position_obj =  vec4f(v_positions[v_idx], 0.0, 1.0);
+    let position_world = ubo.world_from_obj * position_obj;
+    let position_clip = frame.projection_from_world * position_world;
+
+    out.position = position_clip;
+    out.color = v_colors[v_idx] * ubo.color;
+
+    return out;
+}
+
+@fragment
+fn fs_main(in: VertexOut) -> @location(0) vec4f {
+    return in.color;
+}
+
+@fragment
+fn fs_main_picking_layer(in: VertexOut) -> @location(0) vec4u {
+    return vec4u(ubo.picking_layer_object_id, ubo.picking_instance_id);
+}
+
+@fragment
+fn fs_main_outline_mask(in: VertexOut) -> @location(0) vec2u {
+    return ubo.outline_mask;
+}