Support modifying the plot style by introducing a generic framework for overriding components

crates/re_data_ui/src/component_ui_registry.rs

@@ -4,6 +4,8 @@ use re_query::ComponentWithInstances;
 use re_types::external::arrow2::array::Utf8Array;
 use re_viewer_context::{ComponentUiRegistry, UiVerbosity, ViewerContext};
 
+use crate::editors::register_editors;
+
 use super::EntityDataUi;
 
 pub fn create_component_ui_registry() -> ComponentUiRegistry {
@@ -29,6 +31,8 @@ pub fn create_component_ui_registry() -> ComponentUiRegistry {
 
     add_to_registry::<re_types::blueprint::components::IncludedQueries>(&mut registry);
 
+    register_editors(&mut registry);
+
     registry
 }
 

crates/re_data_ui/src/editors.rs

@@ -0,0 +1,221 @@
+// TODO(jleibs): Turn this into a trait
+
+use egui::NumExt as _;
+use re_data_store::{DataStore, LatestAtQuery};
+use re_log_types::EntityPath;
+use re_query::ComponentWithInstances;
+use re_types::{
+    components::{Color, Radius, ScalarScattering, Text},
+    Component, Loggable,
+};
+use re_viewer_context::{UiVerbosity, ViewerContext};
+
+#[allow(clippy::too_many_arguments)]
+fn edit_color_ui(
+    ctx: &ViewerContext<'_>,
+    ui: &mut egui::Ui,
+    _verbosity: UiVerbosity,
+    query: &LatestAtQuery,
+    store: &DataStore,
+    entity_path: &EntityPath,
+    override_path: &EntityPath,
+    component: &ComponentWithInstances,
+    instance_key: &re_types::components::InstanceKey,
+) {
+    let current_color = component
+        .lookup::<Color>(instance_key)
+        .ok()
+        .unwrap_or_else(|| default_color(ctx, query, store, entity_path));
+
+    let [r, g, b, a] = current_color.to_array();
+    let current_color = egui::Color32::from_rgba_unmultiplied(r, g, b, a);
+    let mut edit_color = current_color;
+
+    egui::color_picker::color_edit_button_srgba(
+        ui,
+        &mut edit_color,
+        egui::color_picker::Alpha::Opaque,
+    );
+
+    if edit_color != current_color {
+        let [r, g, b, a] = edit_color.to_array();
+        let new_color = Color::from_unmultiplied_rgba(r, g, b, a);
+
+        ctx.save_blueprint_component(override_path, new_color);
+    }
+}
+
+#[inline]
+fn default_color(
+    _ctx: &ViewerContext<'_>,
+    _query: &LatestAtQuery,
+    _store: &DataStore,
+    _entity_path: &EntityPath,
+) -> Color {
+    Color::from_rgb(255, 255, 255)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn edit_text_ui(
+    ctx: &ViewerContext<'_>,
+    ui: &mut egui::Ui,
+    _verbosity: UiVerbosity,
+    query: &LatestAtQuery,
+    store: &DataStore,
+    entity_path: &EntityPath,
+    override_path: &EntityPath,
+    component: &ComponentWithInstances,
+    instance_key: &re_types::components::InstanceKey,
+) {
+    let current_text = component
+        .lookup::<Text>(instance_key)
+        .ok()
+        .unwrap_or_else(|| default_text(ctx, query, store, entity_path));
+
+    let current_text = current_text.to_string();
+    let mut edit_text = current_text.clone();
+
+    egui::TextEdit::singleline(&mut edit_text).show(ui);
+
+    if edit_text != current_text {
+        let new_text = Text::from(edit_text);
+
+        ctx.save_blueprint_component(override_path, new_text);
+    }
+}
+
+#[inline]
+fn default_text(
+    _ctx: &ViewerContext<'_>,
+    _query: &LatestAtQuery,
+    _store: &DataStore,
+    entity_path: &EntityPath,
+) -> Text {
+    Text::from(entity_path.to_string())
+}
+
+#[allow(clippy::too_many_arguments)]
+fn edit_scatter_ui(
+    ctx: &ViewerContext<'_>,
+    ui: &mut egui::Ui,
+    _verbosity: UiVerbosity,
+    query: &LatestAtQuery,
+    store: &DataStore,
+    entity_path: &EntityPath,
+    override_path: &EntityPath,
+    component: &ComponentWithInstances,
+    instance_key: &re_types::components::InstanceKey,
+) {
+    let current_scatter = component
+        .lookup::<ScalarScattering>(instance_key)
+        .ok()
+        .unwrap_or_else(|| default_scatter(ctx, query, store, entity_path));
+
+    let current_scatter = current_scatter.0;
+    let mut edit_scatter = current_scatter;
+
+    let scattered_text = if current_scatter { "Scattered" } else { "Line" };
+
+    egui::ComboBox::from_id_source("scatter")
+        .selected_text(scattered_text)
+        .show_ui(ui, |ui| {
+            ui.style_mut().wrap = Some(false);
+            ui.selectable_value(&mut edit_scatter, false, "Line");
+            ui.selectable_value(&mut edit_scatter, true, "Scattered");
+        });
+
+    if edit_scatter != current_scatter {
+        let new_scatter = ScalarScattering::from(edit_scatter);
+
+        ctx.save_blueprint_component(override_path, new_scatter);
+    }
+}
+
+#[inline]
+fn default_scatter(
+    _ctx: &ViewerContext<'_>,
+    _query: &LatestAtQuery,
+    _store: &DataStore,
+    _entity_path: &EntityPath,
+) -> ScalarScattering {
+    ScalarScattering::from(false)
+}
+
+#[allow(clippy::too_many_arguments)]
+fn edit_radius_ui(
+    ctx: &ViewerContext<'_>,
+    ui: &mut egui::Ui,
+    _verbosity: UiVerbosity,
+    query: &LatestAtQuery,
+    store: &DataStore,
+    entity_path: &EntityPath,
+    override_path: &EntityPath,
+    component: &ComponentWithInstances,
+    instance_key: &re_types::components::InstanceKey,
+) {
+    let current_radius = component
+        .lookup::<Radius>(instance_key)
+        .ok()
+        .unwrap_or_else(|| default_radius(ctx, query, store, entity_path));
+
+    let current_radius = current_radius.0;
+    let mut edit_radius = current_radius;
+
+    let speed = (current_radius * 0.01).at_least(0.001);
+
+    ui.add(
+        egui::DragValue::new(&mut edit_radius)
+            .clamp_range(0.0..=f64::INFINITY)
+            .speed(speed),
+    );
+
+    if edit_radius != current_radius {
+        let new_radius = Radius::from(edit_radius);
+
+        ctx.save_blueprint_component(override_path, new_radius);
+    }
+}
+
+#[inline]
+fn default_radius(
+    _ctx: &ViewerContext<'_>,
+    _query: &LatestAtQuery,
+    _store: &DataStore,
+    _entity_path: &EntityPath,
+) -> Radius {
+    Radius::from(1.0)
+}
+
+fn register_editor<'a, C: Component + Loggable + 'static>(
+    registry: &mut re_viewer_context::ComponentUiRegistry,
+    default: fn(&ViewerContext<'_>, &LatestAtQuery, &DataStore, &EntityPath) -> C,
+    edit: fn(
+        &ViewerContext<'_>,
+        &mut egui::Ui,
+        UiVerbosity,
+        &LatestAtQuery,
+        &DataStore,
+        &EntityPath,
+        &EntityPath,
+        &ComponentWithInstances,
+        &re_types::components::InstanceKey,
+    ),
+) where
+    C: Into<::std::borrow::Cow<'a, C>>,
+{
+    registry.add_editor(
+        C::name(),
+        Box::new(move |ctx, query, store, entity_path| {
+            let c = default(ctx, query, store, entity_path);
+            [c].into()
+        }),
+        Box::new(edit),
+    );
+}
+
+pub fn register_editors(registry: &mut re_viewer_context::ComponentUiRegistry) {
+    register_editor::<Color>(registry, default_color, edit_color_ui);
+    register_editor::<Text>(registry, default_text, edit_text_ui);
+    register_editor::<ScalarScattering>(registry, default_scatter, edit_scatter_ui);
+    register_editor::<Radius>(registry, default_radius, edit_radius_ui);
+}

crates/re_data_ui/src/lib.rs

@@ -14,6 +14,7 @@ mod component;
 mod component_path;
 mod component_ui_registry;
 mod data;
+mod editors;
 mod entity_db;
 mod entity_path;
 mod image;
@@ -30,6 +31,7 @@ pub use crate::image::{
     show_zoomed_image_region, show_zoomed_image_region_area_outline,
     tensor_summary_ui_grid_contents,
 };
+pub use component::EntityComponentWithInstances;
 pub use component_ui_registry::{add_to_registry, create_component_ui_registry};
 pub use image_meaning::image_meaning_for_entity;
 
@@ -50,7 +52,7 @@ pub fn ui_visible_components<'a>(
 }
 
 /// Show this component in the UI.
-fn is_component_visible_in_ui(component_name: &ComponentName) -> bool {
+pub fn is_component_visible_in_ui(component_name: &ComponentName) -> bool {
     const HIDDEN_COMPONENTS: &[&str] = &["rerun.components.InstanceKey"];
     !HIDDEN_COMPONENTS.contains(&component_name.as_ref())
 }

crates/re_space_view/Cargo.toml

@@ -25,6 +25,7 @@ re_tracing.workspace = true
 re_types_core.workspace = true
 re_viewer_context.workspace = true
 
+ahash.workspace = true
 egui.workspace = true
 itertools.workspace = true
 nohash-hasher.workspace = true

crates/re_space_view/src/data_query_blueprint.rs

@@ -1,3 +1,4 @@
+use ahash::HashMap;
 use nohash_hasher::IntMap;
 use slotmap::SlotMap;
 use smallvec::SmallVec;
@@ -7,9 +8,9 @@ use re_entity_db::{
     EntityPropertyMap, EntityTree,
 };
 use re_log_types::{
-    path::RuleEffect, DataRow, EntityPath, EntityPathFilter, EntityPathRule, RowId,
+    path::RuleEffect, DataRow, EntityPath, EntityPathFilter, EntityPathRule, RowId, StoreKind,
 };
-use re_types_core::archetypes::Clear;
+use re_types_core::{archetypes::Clear, ComponentName};
 use re_viewer_context::{
     blueprint_timepoint_for_writes, DataQueryId, DataQueryResult, DataResult, DataResultHandle,
     DataResultNode, DataResultTree, IndicatorMatchingEntities, PerVisualizer, PropertyOverrides,
@@ -371,6 +372,9 @@ impl DataQueryPropertyResolver<'_> {
             }
         }
 
+        // TODO(jleibs): Should pass through an initial `ComponentOverrides` here
+        // if we were to support incrementally inheriting overrides from parent
+        // contexts such as the `SpaceView` or `Container`.
         EntityOverrideContext {
             root,
             individual: self.resolve_entity_overrides_for_path(
@@ -424,6 +428,8 @@ impl DataQueryPropertyResolver<'_> {
     /// with individual overrides at the leafs.
     fn update_overrides_recursive(
         &self,
+        ctx: &StoreContext<'_>,
+        query: &LatestAtQuery,
         query_result: &mut DataQueryResult,
         override_context: &EntityOverrideContext,
         accumulated: &EntityProperties,
@@ -445,6 +451,7 @@ impl DataQueryPropertyResolver<'_> {
                     node.data_result.property_overrides = Some(PropertyOverrides {
                         individual_properties: overridden_properties.cloned(),
                         accumulated_properties: accumulated_properties.clone(),
+                        component_overrides: Default::default(),
                         override_path: self
                             .recursive_override_root
                             .join(&node.data_result.entity_path),
@@ -462,12 +469,36 @@ impl DataQueryPropertyResolver<'_> {
                         accumulated.clone()
                     };
 
+                    let override_path = self
+                        .individual_override_root
+                        .join(&node.data_result.entity_path);
+
+                    let mut component_overrides: HashMap<ComponentName, (StoreKind, EntityPath)> =
+                        Default::default();
+
+                    if let Some(override_subtree) = ctx.blueprint.tree().subtree(&override_path) {
+                        for component in override_subtree.entity.components.keys() {
+                            if let Some(component_data) = ctx
+                                .blueprint
+                                .store()
+                                .latest_at(query, &override_path, *component, &[*component])
+                                .and_then(|result| result.2[0].clone())
+                            {
+                                if !component_data.is_empty() {
+                                    component_overrides.insert(
+                                        *component,
+                                        (StoreKind::Blueprint, override_path.clone()),
+                                    );
+                                }
+                            }
+                        }
+                    }
+
                     node.data_result.property_overrides = Some(PropertyOverrides {
                         individual_properties: overridden_properties.cloned(),
                         accumulated_properties: accumulated_properties.clone(),
-                        override_path: self
-                            .individual_override_root
-                            .join(&node.data_result.entity_path),
+                        component_overrides,
+                        override_path,
                     });
 
                     None
@@ -476,6 +507,8 @@ impl DataQueryPropertyResolver<'_> {
         {
             for child in child_handles {
                 self.update_overrides_recursive(
+                    ctx,
+                    query,
                     query_result,
                     override_context,
                     &accumulated,
@@ -499,6 +532,8 @@ impl<'a> PropertyResolver for DataQueryPropertyResolver<'a> {
 
         if let Some(root) = query_result.tree.root_handle() {
             self.update_overrides_recursive(
+                ctx,
+                query,
                 query_result,
                 &entity_overrides,
                 &entity_overrides.root,