From 95b2530ba4931e2e6f7622aaddaf9cd9e812172f Mon Sep 17 00:00:00 2001 From: Clement Rey Date: Thu, 18 Apr 2024 16:54:12 +0200 Subject: [PATCH] wip --- Cargo.lock | 7 +- crates/re_entity_db/Cargo.toml | 1 - crates/re_entity_db/src/entity_db.rs | 14 +- crates/re_entity_db/src/lib.rs | 3 +- crates/re_query_cache/src/range/results.rs | 2 + crates/re_space_view/Cargo.toml | 2 +- crates/re_space_view/src/space_view.rs | 4 +- .../re_space_view/src/space_view_contents.rs | 2 +- crates/re_space_view/src/sub_archetypes.rs | 7 +- crates/re_space_view/src/visual_time_range.rs | 2 +- crates/re_space_view_dataframe/Cargo.toml | 2 + .../src/space_view_class.rs | 11 +- .../src/contexts/transform_context.rs | 2 +- crates/re_space_view_spatial/src/lib.rs | 3 +- .../src/visualizers/assets3d.rs | 213 +++++--- .../src/visualizers/entity_iterator.rs | 106 +--- .../src/visualizers/images.rs | 494 +++++++++++------- .../src/visualizers/meshes.rs | 315 +++++++---- .../src/visualizers/results_ext.rs | 10 +- .../src/visualizer_system.rs | 49 +- crates/re_viewer/src/lib.rs | 4 +- crates/re_viewer_context/Cargo.toml | 1 - crates/re_viewer_context/src/annotations.rs | 50 +- crates/re_viewer_context/src/item.rs | 33 +- .../re_viewer_context/src/selection_state.rs | 2 +- .../re_viewer_context/src/space_view/mod.rs | 2 +- crates/re_viewport/src/container.rs | 4 +- crates/re_viewport/src/viewport_blueprint.rs | 4 +- examples/rust/custom_space_view/Cargo.toml | 1 + .../color_coordinates_visualizer_system.rs | 67 ++- 30 files changed, 802 insertions(+), 615 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 462431c67368..b6c34cc5c7ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1406,6 +1406,7 @@ version = "0.16.0-alpha.1+dev" dependencies = [ "mimalloc", "re_crash_handler", + "re_query2", "re_sdk_comms", "re_viewer", "tokio", @@ -4351,7 +4352,6 @@ dependencies = [ "re_log_encoding", "re_log_types", "re_query", - "re_query2", "re_query_cache", "re_smart_channel", "re_tracing", @@ -4718,7 +4718,7 @@ dependencies = [ "re_log", "re_log_types", "re_query", - "re_query2", + "re_query_cache", "re_tracing", "re_types", "re_types_core", @@ -4756,8 +4756,10 @@ dependencies = [ "re_entity_db", "re_log_types", "re_query", + "re_query_cache", "re_renderer", "re_tracing", + "re_types_core", "re_ui", "re_viewer_context", ] @@ -5151,7 +5153,6 @@ dependencies = [ "re_log", "re_log_types", "re_query", - "re_query2", "re_query_cache", "re_renderer", "re_smart_channel", diff --git a/crates/re_entity_db/Cargo.toml b/crates/re_entity_db/Cargo.toml index a2715047f6eb..1ff369ceee67 100644 --- a/crates/re_entity_db/Cargo.toml +++ b/crates/re_entity_db/Cargo.toml @@ -31,7 +31,6 @@ re_log.workspace = true re_log_encoding = { workspace = true, features = ["decoder"] } re_log_types.workspace = true re_query.workspace = true -re_query2.workspace = true re_query_cache = { workspace = true, features = ["to_archetype"] } re_smart_channel.workspace = true re_tracing.workspace = true diff --git a/crates/re_entity_db/src/entity_db.rs b/crates/re_entity_db/src/entity_db.rs index 894b9056a6a0..d9f27bad2b31 100644 --- a/crates/re_entity_db/src/entity_db.rs +++ b/crates/re_entity_db/src/entity_db.rs @@ -12,7 +12,7 @@ use re_log_types::{ LogMsg, RowId, SetStoreInfo, StoreId, StoreInfo, StoreKind, TimePoint, TimeRange, TimeRangeF, Timeline, }; -use re_query2::PromiseResult; +use re_query_cache::PromiseResult; use re_types_core::{components::InstanceKey, Archetype, Loggable}; use crate::{ClearCascade, CompactedStoreEvents, Error, TimesPerTimeline}; @@ -117,7 +117,7 @@ pub struct EntityDb { data_store: DataStore, /// The active promise resolver for this DB. - resolver: re_query2::PromiseResolver, + resolver: re_query_cache::PromiseResolver, /// Query caches for the data in [`Self::data_store`]. query_caches: re_query_cache::Caches, @@ -142,7 +142,7 @@ impl EntityDb { times_per_timeline: Default::default(), tree: crate::EntityTree::root(), data_store, - resolver: re_query2::PromiseResolver::default(), + resolver: re_query_cache::PromiseResolver::default(), query_caches, stats: IngestionStatistics::new(store_id), } @@ -197,7 +197,7 @@ impl EntityDb { } #[inline] - pub fn resolver(&self) -> &re_query2::PromiseResolver { + pub fn resolver(&self) -> &re_query_cache::PromiseResolver { &self.resolver } @@ -207,7 +207,7 @@ impl EntityDb { &self, entity_path: &EntityPath, query: &re_data_store::LatestAtQuery, - ) -> PromiseResult> + ) -> PromiseResult> where re_query_cache::CachedLatestAtResults: re_query_cache::ToArchetype, { @@ -229,7 +229,9 @@ impl EntityDb { } PromiseResult::Error(err) } - PromiseResult::Ready(arch) => PromiseResult::Ready(Some(arch)), + PromiseResult::Ready(arch) => { + PromiseResult::Ready(Some((results.compound_index, arch))) + } } } diff --git a/crates/re_entity_db/src/lib.rs b/crates/re_entity_db/src/lib.rs index 92389cc07aeb..e431475ea747 100644 --- a/crates/re_entity_db/src/lib.rs +++ b/crates/re_entity_db/src/lib.rs @@ -34,7 +34,7 @@ pub(crate) use self::entity_tree::{ClearCascade, CompactedStoreEvents}; use re_log_types::DataTableError; pub use re_log_types::{EntityPath, EntityPathPart, TimeInt, Timeline}; -pub use re_query2::{ExtraQueryHistory, VisibleHistory, VisibleHistoryBoundary}; +pub use re_query_cache::{ExtraQueryHistory, VisibleHistory, VisibleHistoryBoundary}; #[cfg(feature = "serde")] pub use blueprint::components::EntityPropertiesComponent; @@ -43,7 +43,6 @@ pub use editable_auto_value::EditableAutoValue; pub mod external { pub use re_data_store; - pub use re_query2; pub use re_query_cache; } diff --git a/crates/re_query_cache/src/range/results.rs b/crates/re_query_cache/src/range/results.rs index 15482b7b3f65..ce427239cb11 100644 --- a/crates/re_query_cache/src/range/results.rs +++ b/crates/re_query_cache/src/range/results.rs @@ -144,6 +144,8 @@ impl From>> for CachedRangeComponen } } +// TODO: I see no reason why we couldnt put the query range in there, simplifying the life of +// downstream consumers. #[derive(Debug)] pub struct CachedRangeData<'a, T> { // TODO(Amanieu/parking_lot#289): we need two distinct mapped guards because it's diff --git a/crates/re_space_view/Cargo.toml b/crates/re_space_view/Cargo.toml index ef2641b90a71..c40c9293d569 100644 --- a/crates/re_space_view/Cargo.toml +++ b/crates/re_space_view/Cargo.toml @@ -24,7 +24,7 @@ re_log_types.workspace = true re_data_store.workspace = true re_entity_db.workspace = true re_query.workspace = true -re_query2.workspace = true +re_query_cache.workspace = true re_tracing.workspace = true re_types_core.workspace = true re_types.workspace = true diff --git a/crates/re_space_view/src/space_view.rs b/crates/re_space_view/src/space_view.rs index a590f96203c7..cd9b86fa22d7 100644 --- a/crates/re_space_view/src/space_view.rs +++ b/crates/re_space_view/src/space_view.rs @@ -1,6 +1,6 @@ use itertools::{FoldWhile, Itertools}; use nohash_hasher::IntMap; -use re_entity_db::external::re_query2::PromiseResult; +use re_entity_db::external::re_query_cache::PromiseResult; use crate::SpaceViewContents; use re_data_store::LatestAtQuery; @@ -127,7 +127,7 @@ impl SpaceViewBlueprint { // TODO(#5607): what should happen if the promise is still pending? None } - PromiseResult::Ready(arch) => arch, + PromiseResult::Ready(arch) => arch.map(|(_, arch)| arch), PromiseResult::Error(err) => { if cfg!(debug_assertions) { re_log::error!("Failed to load SpaceView blueprint: {err}."); diff --git a/crates/re_space_view/src/space_view_contents.rs b/crates/re_space_view/src/space_view_contents.rs index 4ed247703593..2b4e94cd853e 100644 --- a/crates/re_space_view/src/space_view_contents.rs +++ b/crates/re_space_view/src/space_view_contents.rs @@ -3,7 +3,7 @@ use slotmap::SlotMap; use smallvec::SmallVec; use re_entity_db::{ - external::{re_data_store::LatestAtQuery, re_query2::PromiseResult}, + external::{re_data_store::LatestAtQuery, re_query_cache::PromiseResult}, EntityDb, EntityProperties, EntityPropertiesComponent, EntityPropertyMap, EntityTree, }; use re_log_types::{ diff --git a/crates/re_space_view/src/sub_archetypes.rs b/crates/re_space_view/src/sub_archetypes.rs index ac58a457c598..ab2e94258cc8 100644 --- a/crates/re_space_view/src/sub_archetypes.rs +++ b/crates/re_space_view/src/sub_archetypes.rs @@ -32,7 +32,12 @@ where CachedLatestAtResults: ToArchetype, { let path = entity_path_for_space_view_sub_archetype::(space_view_id, blueprint_db.tree()); - (blueprint_db.latest_at_archetype(&path, query), path) + ( + blueprint_db + .latest_at_archetype(&path, query) + .map(|res| res.map(|(_, arch)| arch)), + path, + ) } pub fn query_space_view_sub_archetype_or_default( diff --git a/crates/re_space_view/src/visual_time_range.rs b/crates/re_space_view/src/visual_time_range.rs index cfa4590fbd06..2c2af59f7363 100644 --- a/crates/re_space_view/src/visual_time_range.rs +++ b/crates/re_space_view/src/visual_time_range.rs @@ -8,7 +8,7 @@ //! to reduce the amount of changes in code that is likely to be refactored soon anyways. use re_log_types::TimeRange; -use re_query2::{ExtraQueryHistory, VisibleHistory, VisibleHistoryBoundary}; +use re_query_cache::{ExtraQueryHistory, VisibleHistory, VisibleHistoryBoundary}; use re_types::blueprint::{ components::VisibleTimeRange, datatypes::{VisibleTimeRangeBoundary, VisibleTimeRangeBoundaryKind}, diff --git a/crates/re_space_view_dataframe/Cargo.toml b/crates/re_space_view_dataframe/Cargo.toml index 0150a8783f2a..88768ab35399 100644 --- a/crates/re_space_view_dataframe/Cargo.toml +++ b/crates/re_space_view_dataframe/Cargo.toml @@ -21,8 +21,10 @@ re_data_ui.workspace = true re_entity_db.workspace = true re_log_types.workspace = true re_query.workspace = true +re_query_cache.workspace = true re_renderer.workspace = true re_tracing.workspace = true +re_types_core.workspace = true re_ui.workspace = true re_viewer_context.workspace = true diff --git a/crates/re_space_view_dataframe/src/space_view_class.rs b/crates/re_space_view_dataframe/src/space_view_class.rs index 06b619a9bc55..448ee5e7cb8f 100644 --- a/crates/re_space_view_dataframe/src/space_view_class.rs +++ b/crates/re_space_view_dataframe/src/space_view_class.rs @@ -6,7 +6,7 @@ use re_data_store::{DataStore, LatestAtQuery}; use re_data_ui::item_ui::instance_path_button; use re_entity_db::{EntityProperties, InstancePath}; use re_log_types::{EntityPath, Timeline}; -use re_query::get_component_with_instances; +use re_types_core::components::InstanceKey; use re_viewer_context::{ SpaceViewClass, SpaceViewClassIdentifier, SpaceViewClassRegistryError, SpaceViewState, SpaceViewSystemExecutionError, SystemExecutionOutput, UiVerbosity, ViewQuery, ViewerContext, @@ -243,9 +243,12 @@ fn sorted_instance_paths_for<'a>( .into_iter() .filter(|comp| !comp.is_indicator_component()) .flat_map(|comp| { - get_component_with_instances(store, latest_at_query, entity_path, comp) - .map(|(_, _, comp_inst)| comp_inst.instance_keys()) - .unwrap_or_default() + let num_instances = store + .latest_at(latest_at_query, entity_path, comp, &[comp]) + .map_or(0, |(_, _, cells)| { + cells[0].as_ref().map_or(0, |cell| cell.num_instances()) + }); + (0..num_instances).map(|i| InstanceKey(i as _)) }) .filter(|instance_key| !instance_key.is_splat()) .collect::>() // dedup and sort diff --git a/crates/re_space_view_spatial/src/contexts/transform_context.rs b/crates/re_space_view_spatial/src/contexts/transform_context.rs index 24b6131ebdc4..8a17eb0f7b54 100644 --- a/crates/re_space_view_spatial/src/contexts/transform_context.rs +++ b/crates/re_space_view_spatial/src/contexts/transform_context.rs @@ -310,7 +310,7 @@ fn get_cached_pinhole( .latest_at_archetype::(entity_path, query) .ok() .flatten() - .map(|arch| { + .map(|(_, arch)| { ( arch.image_from_camera, arch.camera_xyz.unwrap_or(ViewCoordinates::RDF), diff --git a/crates/re_space_view_spatial/src/lib.rs b/crates/re_space_view_spatial/src/lib.rs index d379ed7c70f8..6eec44410cee 100644 --- a/crates/re_space_view_spatial/src/lib.rs +++ b/crates/re_space_view_spatial/src/lib.rs @@ -51,8 +51,7 @@ fn resolution_from_tensor( /// Utility for querying a pinhole archetype instance. /// -/// TODO(andreas): It should be possible to convert [`re_query::ArchetypeView`] to its corresponding Archetype for situations like this. -/// TODO(andreas): This is duplicated into `re_viewport` +// TODO(andreas): This is duplicated into `re_viewport` fn query_pinhole( entity_db: &re_entity_db::EntityDb, query: &re_data_store::LatestAtQuery, diff --git a/crates/re_space_view_spatial/src/visualizers/assets3d.rs b/crates/re_space_view_spatial/src/visualizers/assets3d.rs index 45d201d481a1..e43f5c14fb07 100644 --- a/crates/re_space_view_spatial/src/visualizers/assets3d.rs +++ b/crates/re_space_view_spatial/src/visualizers/assets3d.rs @@ -1,5 +1,6 @@ use re_entity_db::EntityPath; -use re_query::{ArchetypeView, QueryError}; +use re_log_types::{RowId, TimeInt}; +use re_query_cache::{range_zip_1x2, CachedResults}; use re_renderer::renderer::MeshInstance; use re_types::{ archetypes::Asset3D, @@ -11,10 +12,7 @@ use re_viewer_context::{ VisualizerSystem, }; -use super::{ - entity_iterator::process_archetype_views, filter_visualizable_3d_entities, - SpatialViewVisualizerData, -}; +use super::{filter_visualizable_3d_entities, SpatialViewVisualizerData}; use crate::{ contexts::{EntityDepthOffsets, SpatialSceneEntityContext}, instance_hash_conversions::picking_layer_id_from_instance_path_hash, @@ -32,70 +30,80 @@ impl Default for Asset3DVisualizer { } } +struct Asset3DComponentData<'a> { + index: (TimeInt, RowId), + + blob: &'a Blob, + media_type: Option<&'a MediaType>, + transform: Option<&'a OutOfTreeTransform3D>, +} + +// NOTE: Do not put profile scopes in these methods. They are called for all entities and all +// timestamps within a time range -- it's _a lot_. impl Asset3DVisualizer { - fn process_arch_view( + fn process_data<'a>( &mut self, ctx: &ViewerContext<'_>, instances: &mut Vec, - arch_view: &ArchetypeView, - ent_path: &EntityPath, + entity_path: &EntityPath, ent_context: &SpatialSceneEntityContext<'_>, - ) -> Result<(), QueryError> { - let entity_from_pose = arch_view.raw_optional_mono_component::()?; - - let media_type = arch_view.raw_optional_mono_component::()?; - - let mesh = Asset3D { - blob: arch_view.required_mono_component::()?, - media_type: media_type.clone(), - // NOTE: Don't even try to cache the transform! - transform: None, - }; - - let primary_row_id = arch_view.primary_row_id(); - let picking_instance_hash = re_entity_db::InstancePathHash::entity_splat(ent_path); - let outline_mask_ids = ent_context.highlight.index_outline_mask(InstanceKey::SPLAT); - - // TODO(#3232): this is subtly wrong, the key should actually be a hash of everything that got - // cached, which includes the media type… - let mesh = ctx.cache.entry(|c: &mut MeshCache| { - c.entry( - &ent_path.to_string(), - MeshCacheKey { - versioned_instance_path_hash: picking_instance_hash.versioned(primary_row_id), - media_type, - }, - AnyMesh::Asset(&mesh), - ctx.render_ctx, - ) - }); - - if let Some(mesh) = mesh { - re_tracing::profile_scope!("mesh instances"); - - let world_from_pose = ent_context.world_from_entity - * entity_from_pose.map_or(glam::Affine3A::IDENTITY, |t| t.0.into()); - - instances.extend(mesh.mesh_instances.iter().map(move |mesh_instance| { - let pose_from_mesh = mesh_instance.world_from_mesh; - let world_from_mesh = world_from_pose * pose_from_mesh; - - MeshInstance { - gpu_mesh: mesh_instance.gpu_mesh.clone(), - world_from_mesh, - outline_mask_ids, - picking_layer_id: picking_layer_id_from_instance_path_hash( - picking_instance_hash, - ), - ..Default::default() - } - })); - - self.0 - .add_bounding_box(ent_path.hash(), mesh.bbox(), world_from_pose); - }; - - Ok(()) + data: impl Iterator>, + ) { + for data in data { + let mesh = Asset3D { + blob: data.blob.clone(), + media_type: data.media_type.cloned(), + + // NOTE: Don't even try to cache the transform! + transform: None, + }; + + let primary_row_id = data.index.1; + let picking_instance_hash = re_entity_db::InstancePathHash::entity_splat(entity_path); + let outline_mask_ids = ent_context.highlight.index_outline_mask(InstanceKey::SPLAT); + + // TODO(#3232): this is subtly wrong, the key should actually be a hash of everything that got + // cached, which includes the media type… + let mesh = ctx.cache.entry(|c: &mut MeshCache| { + c.entry( + &entity_path.to_string(), + MeshCacheKey { + versioned_instance_path_hash: picking_instance_hash + .versioned(primary_row_id), + media_type: data.media_type.cloned(), + }, + AnyMesh::Asset(&mesh), + ctx.render_ctx, + ) + }); + + if let Some(mesh) = mesh { + re_tracing::profile_scope!("mesh instances"); + + let world_from_pose = ent_context.world_from_entity + * data + .transform + .map_or(glam::Affine3A::IDENTITY, |t| t.0.into()); + + instances.extend(mesh.mesh_instances.iter().map(move |mesh_instance| { + let pose_from_mesh = mesh_instance.world_from_mesh; + let world_from_mesh = world_from_pose * pose_from_mesh; + + MeshInstance { + gpu_mesh: mesh_instance.gpu_mesh.clone(), + world_from_mesh, + outline_mask_ids, + picking_layer_id: picking_layer_id_from_instance_path_hash( + picking_instance_hash, + ), + ..Default::default() + } + })); + + self.0 + .add_bounding_box(entity_path.hash(), mesh.bbox(), world_from_pose); + }; + } } } @@ -122,18 +130,85 @@ impl VisualizerSystem for Asset3DVisualizer { fn execute( &mut self, ctx: &ViewerContext<'_>, - query: &ViewQuery<'_>, + view_query: &ViewQuery<'_>, view_ctx: &ViewContextCollection, ) -> Result, SpaceViewSystemExecutionError> { let mut instances = Vec::new(); - process_archetype_views::( + super::entity_iterator::process_archetype::( ctx, - query, + view_query, view_ctx, view_ctx.get::()?.points, - |ctx, ent_path, _ent_props, arch_view, ent_context| { - self.process_arch_view(ctx, &mut instances, &arch_view, ent_path, ent_context) + |ctx, entity_path, _entity_props, spatial_ctx, results| match results { + CachedResults::LatestAt(_query, results) => { + re_tracing::profile_scope!(format!("{entity_path} @ {_query:?}")); + + use crate::visualizers::CachedLatestAtResultsExt as _; + + let resolver = ctx.recording().resolver(); + + let blob = match results.get_dense::(resolver) { + Some(Ok(blobs)) if !blobs.is_empty() => blobs.first().unwrap(), + Some(Err(err)) => return Err(err.into()), + _ => return Ok(()), + }; + + let media_types = results.get_or_empty_dense(resolver)?; + let transforms = results.get_or_empty_dense(resolver)?; + + let data = { + Asset3DComponentData { + index: results.compound_index, + blob, + media_type: media_types.first(), + transform: transforms.first(), + } + }; + + self.process_data( + ctx, + &mut instances, + entity_path, + spatial_ctx, + std::iter::once(data), + ); + Ok(()) + } + + CachedResults::Range(_query, results) => { + re_tracing::profile_scope!(format!("{entity_path} @ {_query:?}")); + + use crate::visualizers::CachedRangeResultsExt as _; + + let resolver = ctx.recording().resolver(); + + let blobs = match results.get_dense::(resolver, _query) { + Some(Ok(blobs)) => blobs, + Some(err @ Err(_)) => err?, + _ => return Ok(()), + }; + + let media_types = results.get_or_empty_dense(resolver, _query)?; + let transforms = results.get_or_empty_dense(resolver, _query)?; + + let data = range_zip_1x2( + blobs.range_indexed(_query.range()), + media_types.range_indexed(_query.range()), + transforms.range_indexed(_query.range()), + ) + .filter_map(|(&index, blobs, media_types, transforms)| { + blobs.first().map(|blob| Asset3DComponentData { + index, + blob, + media_type: media_types.and_then(|media_types| media_types.first()), + transform: transforms.and_then(|transforms| transforms.first()), + }) + }); + + self.process_data(ctx, &mut instances, entity_path, spatial_ctx, data); + Ok(()) + } }, )?; diff --git a/crates/re_space_view_spatial/src/visualizers/entity_iterator.rs b/crates/re_space_view_spatial/src/visualizers/entity_iterator.rs index f63691dc3753..7b341e5064a6 100644 --- a/crates/re_space_view_spatial/src/visualizers/entity_iterator.rs +++ b/crates/re_space_view_spatial/src/visualizers/entity_iterator.rs @@ -2,7 +2,6 @@ use itertools::Either; use re_data_store::{LatestAtQuery, RangeQuery}; use re_entity_db::{EntityDb, EntityProperties}; use re_log_types::{EntityPath, TimeInt, Timeline}; -use re_query::{ArchetypeView, QueryError}; use re_query_cache::{CachedResults, ExtraQueryHistory}; use re_renderer::DepthOffset; use re_space_view::query_visual_history; @@ -39,6 +38,8 @@ pub fn clamped(values: &[T], clamped_len: usize) -> impl Iterator ) } +// --- Cached APIs --- + pub fn query_archetype_with_history( entity_db: &EntityDb, timeline: &Timeline, @@ -161,106 +162,3 @@ where Ok(()) } - -// --- - -/// Iterates through all entity views for a given archetype. -/// -/// The callback passed in gets passed a long an [`SpatialSceneEntityContext`] which contains -/// various useful information about an entity in the context of the current scene. -pub fn process_archetype_views<'a, System: IdentifiedViewSystem, A, const N: usize, F>( - ctx: &ViewerContext<'_>, - query: &ViewQuery<'_>, - view_ctx: &ViewContextCollection, - default_depth_offset: DepthOffset, - mut fun: F, -) -> Result<(), SpaceViewSystemExecutionError> -where - A: Archetype + 'a, - F: FnMut( - &ViewerContext<'_>, - &EntityPath, - &EntityProperties, - ArchetypeView, - &SpatialSceneEntityContext<'_>, - ) -> Result<(), QueryError>, -{ - let transforms = view_ctx.get::()?; - let depth_offsets = view_ctx.get::()?; - let annotations = view_ctx.get::()?; - let counter = view_ctx.get::()?; - - for data_result in query.iter_visible_data_results(ctx, System::identifier()) { - // The transform that considers pinholes only makes sense if this is a 3D space-view - let world_from_entity = - if view_ctx.space_view_class_identifier() == SpatialSpaceView3D::identifier() { - transforms.reference_from_entity(&data_result.entity_path) - } else { - transforms.reference_from_entity_ignoring_pinhole( - &data_result.entity_path, - ctx.recording(), - &query.latest_at_query(), - ) - }; - - let Some(world_from_entity) = world_from_entity else { - continue; - }; - let entity_context = SpatialSceneEntityContext { - world_from_entity, - depth_offset: *depth_offsets - .per_entity - .get(&data_result.entity_path.hash()) - .unwrap_or(&default_depth_offset), - annotations: annotations.0.find(&data_result.entity_path), - highlight: query - .highlights - .entity_outline_mask(data_result.entity_path.hash()), - space_view_class_identifier: view_ctx.space_view_class_identifier(), - }; - - let extra_history = query_visual_history(ctx, data_result); - // TODO(cmc): We need to use the type defined in the old crate interchangeably with the one - // defined in the new crate. They are exactly the same. - // - // This will be fixed in the next PR - #[allow(unsafe_code, clippy::undocumented_unsafe_blocks)] - let extra_history = unsafe { std::mem::transmute(extra_history) }; - - let result = re_query::query_archetype_with_history::( - ctx.recording_store(), - &query.timeline, - &query.latest_at, - &extra_history, - &data_result.entity_path, - ) - .and_then(|arch_views| { - for arch_view in arch_views { - counter.num_primitives.fetch_add( - arch_view.num_instances(), - std::sync::atomic::Ordering::Relaxed, - ); - - fun( - ctx, - &data_result.entity_path, - data_result.accumulated_properties(), - arch_view, - &entity_context, - )?; - } - Ok(()) - }); - match result { - Ok(_) | Err(QueryError::PrimaryNotFound(_)) => {} - Err(err) => { - re_log::error_once!( - "Unexpected error querying {:?}: {err}", - &data_result.entity_path - ); - } - } - } - - Ok(()) -} diff --git a/crates/re_space_view_spatial/src/visualizers/images.rs b/crates/re_space_view_spatial/src/visualizers/images.rs index 7fc8f9671c0c..693d53312f06 100644 --- a/crates/re_space_view_spatial/src/visualizers/images.rs +++ b/crates/re_space_view_spatial/src/visualizers/images.rs @@ -5,8 +5,8 @@ use itertools::Itertools as _; use nohash_hasher::IntSet; use re_entity_db::{EntityPath, EntityProperties}; -use re_log_types::{EntityPathHash, RowId}; -use re_query::{ArchetypeView, QueryError}; +use re_log_types::{EntityPathHash, RowId, TimeInt}; +use re_query_cache::{range_zip_1x2, CachedResults}; use re_renderer::{ renderer::{DepthCloud, DepthClouds, RectangleOptions, TexturedRect}, Colormap, @@ -16,7 +16,7 @@ use re_types::{ archetypes::{DepthImage, Image, SegmentationImage}, components::{Color, DrawOrder, TensorData, ViewCoordinates}, tensor_data::{DecodedTensor, TensorDataMeaning}, - Archetype as _, ComponentNameSet, + Archetype, ComponentNameSet, }; use re_viewer_context::{ gpu_bridge, ApplicableEntities, DefaultColor, IdentifiedViewSystem, SpaceViewClass, @@ -33,7 +33,7 @@ use crate::{ SpatialSpaceView2D, SpatialSpaceView3D, }; -use super::{entity_iterator::process_archetype_views, SpatialViewVisualizerData}; +use super::SpatialViewVisualizerData; pub struct ViewerImage { /// Path to the image (note image instance ids would refer to pixels!) @@ -64,8 +64,6 @@ fn to_textured_rect( meaning: TensorDataMeaning, multiplicative_tint: egui::Rgba, ) -> Option { - re_tracing::profile_function!(); - let [height, width, _] = tensor.image_height_width_channels()?; let debug_name = ent_path.to_string(); @@ -149,10 +147,18 @@ impl Default for ImageVisualizer { } } +struct ImageComponentData<'a> { + index: (TimeInt, RowId), + + tensor: &'a TensorData, + color: Option<&'a Color>, + draw_order: Option<&'a DrawOrder>, +} + +// NOTE: Do not put profile scopes in these methods. They are called for all entities and all +// timestamps within a time range -- it's _a lot_. impl ImageVisualizer { fn handle_image_layering(&mut self) { - re_tracing::profile_function!(); - // Rebuild the image list, grouped by "shared plane", identified with camera & draw order. let mut image_groups: BTreeMap> = BTreeMap::new(); for image in self.images.drain(..) { @@ -201,60 +207,50 @@ impl ImageVisualizer { } #[allow(clippy::too_many_arguments)] - fn process_image_arch_view( + fn process_image_data<'a>( &mut self, ctx: &ViewerContext<'_>, transforms: &TransformContext, ent_props: &EntityProperties, - arch_view: &ArchetypeView, - ent_path: &EntityPath, + entity_path: &EntityPath, ent_context: &SpatialSceneEntityContext<'_>, - ) -> Result<(), QueryError> { - re_tracing::profile_function!(); + data: impl Iterator>, + ) { + // If this isn't an image, return + // TODO(jleibs): The ArchetypeView should probably do this for us. + if !ctx.recording_store().entity_has_component( + &ctx.current_query().timeline(), + entity_path, + &Image::indicator().name(), + ) { + return; + } // Parent pinhole should only be relevant to 3D views let parent_pinhole_path = if ent_context.space_view_class_identifier == SpatialSpaceView3D::identifier() { - transforms.parent_pinhole(ent_path) + transforms.parent_pinhole(entity_path) } else { None }; - // If this isn't an image, return - // TODO(jleibs): The ArchetypeView should probably do this for us. - if !ctx.recording_store().entity_has_component( - &ctx.current_query().timeline(), - ent_path, - &Image::indicator().name(), - ) { - return Ok(()); - } // Unknown is currently interpreted as "Some Color" in most cases. // TODO(jleibs): Make this more explicit let meaning = TensorDataMeaning::Unknown; - // Instance ids of tensors refer to entries inside the tensor. - for (tensor, color, draw_order) in itertools::izip!( - arch_view.iter_required_component::()?, - arch_view.iter_optional_component::()?, - arch_view.iter_optional_component::()? - ) { - re_tracing::profile_scope!("loop_iter"); - - if !tensor.is_shaped_like_an_image() { - return Ok(()); + for data in data { + if !data.tensor.is_shaped_like_an_image() { + continue; } - let tensor_data_row_id = arch_view.primary_row_id(); - - let tensor = match ctx - .cache - .entry(|c: &mut TensorDecodeCache| c.entry(tensor_data_row_id, tensor.0)) - { + let tensor_data_row_id = data.index.1; + let tensor = match ctx.cache.entry(|c: &mut TensorDecodeCache| { + c.entry(tensor_data_row_id, data.tensor.0.clone()) + }) { Ok(tensor) => tensor, Err(err) => { re_log::warn_once!( - "Encountered problem decoding tensor at path {ent_path}: {err}" + "Encountered problem decoding tensor at path {entity_path}: {err}" ); continue; } @@ -264,11 +260,11 @@ impl ImageVisualizer { .annotations .resolved_class_description(None) .annotation_info() - .color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite); + .color(data.color.map(|c| c.to_array()), DefaultColor::OpaqueWhite); if let Some(textured_rect) = to_textured_rect( ctx, - ent_path, + entity_path, ent_context, tensor_data_row_id, &tensor, @@ -284,129 +280,81 @@ impl ImageVisualizer { || !ent_props.pinhole_image_plane_distance.is_auto() { self.data.add_bounding_box( - ent_path.hash(), + entity_path.hash(), Self::compute_bounding_box(&textured_rect), ent_context.world_from_entity, ); } self.images.push(ViewerImage { - ent_path: ent_path.clone(), + ent_path: entity_path.clone(), tensor, meaning, textured_rect, parent_pinhole: parent_pinhole_path.map(|p| p.hash()), - draw_order: draw_order.unwrap_or(DrawOrder::DEFAULT_IMAGE), + draw_order: data.draw_order.copied().unwrap_or(DrawOrder::DEFAULT_IMAGE), }); } } - - Ok(()) } #[allow(clippy::too_many_arguments)] - fn process_depth_image_arch_view( + fn process_segmentation_image_data<'a>( &mut self, ctx: &ViewerContext<'_>, - depth_clouds: &mut Vec, transforms: &TransformContext, ent_props: &EntityProperties, - arch_view: &ArchetypeView, - ent_path: &EntityPath, + entity_path: &EntityPath, ent_context: &SpatialSceneEntityContext<'_>, - ) -> Result<(), QueryError> { - re_tracing::profile_function!(); - + data: impl Iterator>, + ) { // If this isn't an image, return // TODO(jleibs): The ArchetypeView should probably to this for us. if !ctx.recording_store().entity_has_component( &ctx.current_query().timeline(), - ent_path, - &DepthImage::indicator().name(), + entity_path, + &SegmentationImage::indicator().name(), ) { - return Ok(()); + return; } - let meaning = TensorDataMeaning::Depth; // Parent pinhole should only be relevant to 3D views let parent_pinhole_path = if ent_context.space_view_class_identifier == SpatialSpaceView3D::identifier() { - transforms.parent_pinhole(ent_path) + transforms.parent_pinhole(entity_path) } else { None }; - // Instance ids of tensors refer to entries inside the tensor. - for (tensor, color, draw_order) in itertools::izip!( - arch_view.iter_required_component::()?, - arch_view.iter_optional_component::()?, - arch_view.iter_optional_component::()? - ) { - // NOTE: we ignore the `DepthMeter` component here because we get it from - // `EntityProperties::depth_from_world_scale` instead, which is initialized to the - // same value, but the user may have edited it. - re_tracing::profile_scope!("loop_iter"); + let meaning = TensorDataMeaning::ClassId; - if !tensor.is_shaped_like_an_image() { - return Ok(()); + for data in data { + if !data.tensor.is_shaped_like_an_image() { + continue; } - let tensor_data_row_id = arch_view.primary_row_id(); - - let tensor = match ctx - .cache - .entry(|c: &mut TensorDecodeCache| c.entry(tensor_data_row_id, tensor.0)) - { + let tensor_data_row_id = data.index.1; + let tensor = match ctx.cache.entry(|c: &mut TensorDecodeCache| { + c.entry(tensor_data_row_id, data.tensor.0.clone()) + }) { Ok(tensor) => tensor, Err(err) => { re_log::warn_once!( - "Encountered problem decoding tensor at path {ent_path}: {err}" + "Encountered problem decoding tensor at path {entity_path}: {err}" ); continue; } }; - if *ent_props.backproject_depth { - if let Some(parent_pinhole_path) = transforms.parent_pinhole(ent_path) { - // NOTE: we don't pass in `world_from_obj` because this corresponds to the - // transform of the projection plane, which is of no use to us here. - // What we want are the extrinsics of the depth camera! - match Self::process_entity_view_as_depth_cloud( - ctx, - transforms, - ent_context, - ent_props, - tensor_data_row_id, - &tensor, - ent_path, - parent_pinhole_path, - ) { - Ok(cloud) => { - self.data.add_bounding_box( - ent_path.hash(), - cloud.world_space_bbox(), - glam::Affine3A::IDENTITY, - ); - self.depth_cloud_entities.insert(ent_path.hash()); - depth_clouds.push(cloud); - return Ok(()); - } - Err(err) => { - re_log::warn_once!("{err}"); - } - } - }; - } - let color = ent_context .annotations .resolved_class_description(None) .annotation_info() - .color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite); + .color(data.color.map(|c| c.to_array()), DefaultColor::OpaqueWhite); if let Some(textured_rect) = to_textured_rect( ctx, - ent_path, + entity_path, ent_context, tensor_data_row_id, &tensor, @@ -422,94 +370,118 @@ impl ImageVisualizer { || !ent_props.pinhole_image_plane_distance.is_auto() { self.data.add_bounding_box( - ent_path.hash(), + entity_path.hash(), Self::compute_bounding_box(&textured_rect), ent_context.world_from_entity, ); } self.images.push(ViewerImage { - ent_path: ent_path.clone(), + ent_path: entity_path.clone(), tensor, meaning, textured_rect, parent_pinhole: parent_pinhole_path.map(|p| p.hash()), - draw_order: draw_order.unwrap_or(DrawOrder::DEFAULT_IMAGE), + draw_order: data.draw_order.copied().unwrap_or(DrawOrder::DEFAULT_IMAGE), }); } } - - Ok(()) } #[allow(clippy::too_many_arguments)] - fn process_segmentation_image_arch_view( + fn process_depth_image_data<'a>( &mut self, ctx: &ViewerContext<'_>, + depth_clouds: &mut Vec, transforms: &TransformContext, ent_props: &EntityProperties, - arch_view: &ArchetypeView, - ent_path: &EntityPath, + entity_path: &EntityPath, ent_context: &SpatialSceneEntityContext<'_>, - ) -> Result<(), QueryError> { - re_tracing::profile_function!(); + data: impl Iterator>, + ) { + // If this isn't an image, return + // TODO(jleibs): The ArchetypeView should probably to this for us. + if !ctx.recording_store().entity_has_component( + &ctx.current_query().timeline(), + entity_path, + &DepthImage::indicator().name(), + ) { + return; + } // Parent pinhole should only be relevant to 3D views let parent_pinhole_path = if ent_context.space_view_class_identifier == SpatialSpaceView3D::identifier() { - transforms.parent_pinhole(ent_path) + transforms.parent_pinhole(entity_path) } else { None }; - // If this isn't an image, return - // TODO(jleibs): The ArchetypeView should probably to this for us. - if !ctx.recording_store().entity_has_component( - &ctx.current_query().timeline(), - ent_path, - &SegmentationImage::indicator().name(), - ) { - return Ok(()); - } - - let meaning = TensorDataMeaning::ClassId; + let meaning = TensorDataMeaning::Depth; - // Instance ids of tensors refer to entries inside the tensor. - for (tensor, color, draw_order) in itertools::izip!( - arch_view.iter_required_component::()?, - arch_view.iter_optional_component::()?, - arch_view.iter_optional_component::()? - ) { - re_tracing::profile_scope!("loop_iter"); + for data in data { + // NOTE: we ignore the `DepthMeter` component here because we get it from + // `EntityProperties::depth_from_world_scale` instead, which is initialized to the + // same value, but the user may have edited it. - if !tensor.is_shaped_like_an_image() { - return Ok(()); + if !data.tensor.is_shaped_like_an_image() { + continue; } - let tensor_data_row_id = arch_view.primary_row_id(); - - let tensor = match ctx - .cache - .entry(|c: &mut TensorDecodeCache| c.entry(tensor_data_row_id, tensor.0)) - { + let tensor_data_row_id = data.index.1; + let tensor = match ctx.cache.entry(|c: &mut TensorDecodeCache| { + c.entry(tensor_data_row_id, data.tensor.0.clone()) + }) { Ok(tensor) => tensor, Err(err) => { re_log::warn_once!( - "Encountered problem decoding tensor at path {ent_path}: {err}" + "Encountered problem decoding tensor at path {entity_path}: {err}" ); continue; } }; + if *ent_props.backproject_depth { + if let Some(parent_pinhole_path) = transforms.parent_pinhole(entity_path) { + // NOTE: we don't pass in `world_from_obj` because this corresponds to the + // transform of the projection plane, which is of no use to us here. + // What we want are the extrinsics of the depth camera! + match Self::process_entity_view_as_depth_cloud( + ctx, + transforms, + ent_context, + ent_props, + tensor_data_row_id, + &tensor, + entity_path, + parent_pinhole_path, + ) { + Ok(cloud) => { + self.data.add_bounding_box( + entity_path.hash(), + cloud.world_space_bbox(), + glam::Affine3A::IDENTITY, + ); + self.depth_cloud_entities.insert(entity_path.hash()); + depth_clouds.push(cloud); + return; + } + Err(err) => { + re_log::warn_once!("{err}"); + } + } + }; + } + let color = ent_context .annotations .resolved_class_description(None) .annotation_info() - .color(color.map(|c| c.to_array()), DefaultColor::OpaqueWhite); + .color(data.color.map(|c| c.to_array()), DefaultColor::OpaqueWhite); if let Some(textured_rect) = to_textured_rect( ctx, - ent_path, + entity_path, ent_context, tensor_data_row_id, &tensor, @@ -525,24 +497,22 @@ impl ImageVisualizer { || !ent_props.pinhole_image_plane_distance.is_auto() { self.data.add_bounding_box( - ent_path.hash(), + entity_path.hash(), Self::compute_bounding_box(&textured_rect), ent_context.world_from_entity, ); } self.images.push(ViewerImage { - ent_path: ent_path.clone(), + ent_path: entity_path.clone(), tensor, meaning, textured_rect, parent_pinhole: parent_pinhole_path.map(|p| p.hash()), - draw_order: draw_order.unwrap_or(DrawOrder::DEFAULT_IMAGE), + draw_order: data.draw_order.copied().unwrap_or(DrawOrder::DEFAULT_IMAGE), }); } } - - Ok(()) } #[allow(clippy::too_many_arguments)] @@ -729,67 +699,81 @@ impl VisualizerSystem for ImageVisualizer { fn execute( &mut self, ctx: &ViewerContext<'_>, - query: &ViewQuery<'_>, + view_query: &ViewQuery<'_>, view_ctx: &ViewContextCollection, ) -> Result, SpaceViewSystemExecutionError> { let mut depth_clouds = Vec::new(); - let transforms = view_ctx.get::()?; - - process_archetype_views::( + self.process_image_archetype::( ctx, - query, + view_query, view_ctx, - view_ctx.get::()?.image, - |ctx, ent_path, ent_props, ent_view, ent_context| { - self.process_image_arch_view( + &mut depth_clouds, + |visualizer, + ctx, + _depth_clouds, + transforms, + entity_props, + entity_path, + spatial_ctx, + data| { + visualizer.process_image_data( ctx, transforms, - ent_props, - &ent_view, - ent_path, - ent_context, - ) + entity_props, + entity_path, + spatial_ctx, + data, + ); }, )?; - process_archetype_views::< - ImageVisualizer, - SegmentationImage, - { SegmentationImage::NUM_COMPONENTS }, - _, - >( + self.process_image_archetype::( ctx, - query, + view_query, view_ctx, - view_ctx.get::()?.image, - |ctx, ent_path, ent_props, ent_view, ent_context| { - self.process_segmentation_image_arch_view( + &mut depth_clouds, + |visualizer, + ctx, + _depth_clouds, + transforms, + entity_props, + entity_path, + spatial_ctx, + data| { + visualizer.process_segmentation_image_data( ctx, transforms, - ent_props, - &ent_view, - ent_path, - ent_context, - ) + entity_props, + entity_path, + spatial_ctx, + data, + ); }, )?; - process_archetype_views::( + self.process_image_archetype::( ctx, - query, + view_query, view_ctx, - view_ctx.get::()?.image, - |ctx, ent_path, ent_props, ent_view, ent_context| { - self.process_depth_image_arch_view( + &mut depth_clouds, + |visualizer, + ctx, + depth_clouds, + transforms, + entity_props, + entity_path, + spatial_ctx, + data| { + visualizer.process_depth_image_data( ctx, - &mut depth_clouds, + depth_clouds, transforms, - ent_props, - &ent_view, - ent_path, - ent_context, - ) + entity_props, + entity_path, + spatial_ctx, + data, + ); }, )?; @@ -839,3 +823,121 @@ impl VisualizerSystem for ImageVisualizer { self } } + +impl ImageVisualizer { + fn process_image_archetype( + &mut self, + ctx: &ViewerContext<'_>, + view_query: &ViewQuery<'_>, + view_ctx: &ViewContextCollection, + depth_clouds: &mut Vec, + mut f: F, + ) -> Result<(), SpaceViewSystemExecutionError> + where + F: FnMut( + &mut ImageVisualizer, + &ViewerContext<'_>, + &mut Vec, + &TransformContext, + &EntityProperties, + &EntityPath, + &SpatialSceneEntityContext<'_>, + &mut dyn Iterator>, + ), + { + let transforms = view_ctx.get::()?; + + super::entity_iterator::process_archetype::( + ctx, + view_query, + view_ctx, + view_ctx.get::()?.image, + |ctx, entity_path, entity_props, spatial_ctx, results| match results { + CachedResults::LatestAt(_query, results) => { + re_tracing::profile_scope!(format!("{entity_path} @ {_query:?}")); + + use crate::visualizers::CachedLatestAtResultsExt as _; + + let resolver = ctx.recording().resolver(); + + let tensor = match results.get_dense::(resolver) { + Some(Ok(tensors)) if !tensors.is_empty() => tensors.first().unwrap(), + Some(Err(err)) => return Err(err.into()), + _ => return Ok(()), + }; + + let colors = results.get_or_empty_dense(resolver)?; + let draw_orders = results.get_or_empty_dense(resolver)?; + + let data = { + ImageComponentData { + index: results.compound_index, + tensor, + color: colors.first(), + draw_order: draw_orders.first(), + } + }; + + f( + self, + ctx, + depth_clouds, + transforms, + entity_props, + entity_path, + spatial_ctx, + &mut std::iter::once(data), + ); + + Ok(()) + } + + CachedResults::Range(_query, results) => { + re_tracing::profile_scope!(format!("{entity_path} @ {_query:?}")); + + use crate::visualizers::CachedRangeResultsExt as _; + + let resolver = ctx.recording().resolver(); + + let tensors = match results.get_dense::(resolver, _query) { + Some(Ok(tensors)) => tensors, + Some(err @ Err(_)) => err?, + _ => return Ok(()), + }; + + let colors = results.get_or_empty_dense(resolver, _query)?; + let draw_orders = results.get_or_empty_dense(resolver, _query)?; + + let mut data = range_zip_1x2( + tensors.range_indexed(_query.range()), + draw_orders.range_indexed(_query.range()), + colors.range_indexed(_query.range()), + ) + .filter_map(|(&index, tensors, draw_orders, colors)| { + tensors.first().map(|tensor| ImageComponentData { + index, + tensor, + color: colors.and_then(|colors| colors.first()), + draw_order: draw_orders.and_then(|draw_orders| draw_orders.first()), + }) + }); + + f( + self, + ctx, + depth_clouds, + transforms, + entity_props, + entity_path, + spatial_ctx, + &mut data, + ); + + Ok(()) + } + }, + )?; + + Ok(()) + } +} diff --git a/crates/re_space_view_spatial/src/visualizers/meshes.rs b/crates/re_space_view_spatial/src/visualizers/meshes.rs index 97f8a15344d7..ea00addc3d19 100644 --- a/crates/re_space_view_spatial/src/visualizers/meshes.rs +++ b/crates/re_space_view_spatial/src/visualizers/meshes.rs @@ -1,10 +1,13 @@ +use itertools::Itertools as _; use re_entity_db::EntityPath; -use re_query::{ArchetypeView, QueryError}; +use re_log_types::{RowId, TimeInt}; +use re_query_cache::{range_zip_1x7, CachedResults}; use re_renderer::renderer::MeshInstance; use re_types::{ archetypes::Mesh3D, components::{ - Color, InstanceKey, Material, MeshProperties, Position3D, TensorData, Texcoord2D, Vector3D, + ClassId, Color, InstanceKey, Material, MeshProperties, Position3D, TensorData, Texcoord2D, + Vector3D, }, }; use re_viewer_context::{ @@ -13,10 +16,6 @@ use re_viewer_context::{ VisualizerSystem, }; -use super::{ - entity_iterator::process_archetype_views, filter_visualizable_3d_entities, - SpatialViewVisualizerData, -}; use crate::{ contexts::{EntityDepthOffsets, SpatialSceneEntityContext}, instance_hash_conversions::picking_layer_id_from_instance_path_hash, @@ -24,6 +23,10 @@ use crate::{ view_kind::SpatialSpaceViewKind, }; +use super::{entity_iterator::clamped, filter_visualizable_3d_entities, SpatialViewVisualizerData}; + +// --- + pub struct Mesh3DVisualizer(SpatialViewVisualizerData); impl Default for Mesh3DVisualizer { @@ -34,116 +37,98 @@ impl Default for Mesh3DVisualizer { } } +struct Mesh3DComponentData<'a> { + index: (TimeInt, RowId), + + vertex_positions: &'a [Position3D], + vertex_normals: &'a [Vector3D], + vertex_colors: &'a [Color], + vertex_texcoords: &'a [Texcoord2D], + + mesh_properties: Option<&'a MeshProperties>, + mesh_material: Option<&'a Material>, + albedo_texture: Option<&'a TensorData>, + + class_ids: &'a [ClassId], +} + +// NOTE: Do not put profile scopes in these methods. They are called for all entities and all +// timestamps within a time range -- it's _a lot_. impl Mesh3DVisualizer { - fn process_arch_view( + fn process_data<'a>( &mut self, ctx: &ViewerContext<'_>, instances: &mut Vec, - arch_view: &ArchetypeView, - ent_path: &EntityPath, + entity_path: &EntityPath, ent_context: &SpatialSceneEntityContext<'_>, - ) -> Result<(), QueryError> { - re_tracing::profile_function!(); + data: impl Iterator>, + ) { + for data in data { + let primary_row_id = data.index.1; + let picking_instance_hash = re_entity_db::InstancePathHash::entity_splat(entity_path); + let outline_mask_ids = ent_context.highlight.index_outline_mask(InstanceKey::SPLAT); - let vertex_positions: Vec<_> = { - re_tracing::profile_scope!("vertex_positions"); - arch_view.iter_required_component::()?.collect() - }; - if vertex_positions.is_empty() { - return Ok(()); - } + let mesh = ctx.cache.entry(|c: &mut MeshCache| { + let key = MeshCacheKey { + versioned_instance_path_hash: picking_instance_hash.versioned(primary_row_id), + media_type: None, + }; - let mesh = { - re_tracing::profile_scope!("collect"); - // NOTE: - // - Per-vertex properties are joined using the cluster key as usual. - // - Per-mesh properties are just treated as a "global var", essentially. - Mesh3D { - vertex_positions, - vertex_normals: if arch_view.has_component::() { - re_tracing::profile_scope!("vertex_normals"); - Some( - arch_view - .iter_optional_component::()? - .map(|comp| comp.unwrap_or(Vector3D::ZERO)) - .collect(), - ) - } else { - None - }, - vertex_colors: if arch_view.has_component::() { - re_tracing::profile_scope!("vertex_colors"); - let fallback = Color::new(0xFFFFFFFF); - Some( - arch_view - .iter_optional_component::()? - .map(|comp| comp.unwrap_or(fallback)) - .collect(), - ) - } else { - None - }, - vertex_texcoords: if arch_view.has_component::() { - re_tracing::profile_scope!("vertex_texcoords"); - Some( - arch_view - .iter_optional_component::()? - .map(|comp| comp.unwrap_or(Texcoord2D::ZERO)) - .collect(), - ) - } else { - None - }, - mesh_properties: arch_view.raw_optional_mono_component::()?, - mesh_material: arch_view.raw_optional_mono_component::()?, - albedo_texture: arch_view.raw_optional_mono_component::()?, - class_ids: None, - } - }; + let vertex_normals = clamped(data.vertex_normals, data.vertex_positions.len()) + .copied() + .collect_vec(); + let vertex_colors = clamped(data.vertex_colors, data.vertex_positions.len()) + .copied() + .collect_vec(); + let vertex_texcoords = clamped(data.vertex_texcoords, data.vertex_positions.len()) + .copied() + .collect_vec(); - let primary_row_id = arch_view.primary_row_id(); - let picking_instance_hash = re_entity_db::InstancePathHash::entity_splat(ent_path); - let outline_mask_ids = ent_context.highlight.index_outline_mask(InstanceKey::SPLAT); + c.entry( + &entity_path.to_string(), + key.clone(), + AnyMesh::Mesh { + mesh: &Mesh3D { + vertex_positions: data.vertex_positions.to_owned(), + mesh_properties: data.mesh_properties.cloned(), + vertex_normals: (!vertex_normals.is_empty()).then_some(vertex_normals), + vertex_colors: (!vertex_colors.is_empty()).then_some(vertex_colors), + vertex_texcoords: (!vertex_texcoords.is_empty()) + .then_some(vertex_texcoords), + mesh_material: data.mesh_material.cloned(), + albedo_texture: data.albedo_texture.cloned(), + class_ids: (!data.class_ids.is_empty()) + .then(|| data.class_ids.to_owned()), + }, + texture_key: re_log_types::hash::Hash64::hash(&key).hash64(), + }, + ctx.render_ctx, + ) + }); - let mesh = ctx.cache.entry(|c: &mut MeshCache| { - let key = MeshCacheKey { - versioned_instance_path_hash: picking_instance_hash.versioned(primary_row_id), - media_type: None, - }; - c.entry( - &ent_path.to_string(), - key.clone(), - AnyMesh::Mesh { - mesh: &mesh, - texture_key: re_log_types::hash::Hash64::hash(&key).hash64(), - }, - ctx.render_ctx, - ) - }); - - if let Some(mesh) = mesh { - re_tracing::profile_scope!("mesh instances"); - - instances.extend(mesh.mesh_instances.iter().map(move |mesh_instance| { - let entity_from_mesh = mesh_instance.world_from_mesh; - let world_from_mesh = ent_context.world_from_entity * entity_from_mesh; - - MeshInstance { - gpu_mesh: mesh_instance.gpu_mesh.clone(), - world_from_mesh, - outline_mask_ids, - picking_layer_id: picking_layer_id_from_instance_path_hash( - picking_instance_hash, - ), - ..Default::default() - } - })); + if let Some(mesh) = mesh { + instances.extend(mesh.mesh_instances.iter().map(move |mesh_instance| { + let entity_from_mesh = mesh_instance.world_from_mesh; + let world_from_mesh = ent_context.world_from_entity * entity_from_mesh; - self.0 - .add_bounding_box(ent_path.hash(), mesh.bbox(), ent_context.world_from_entity); - }; + MeshInstance { + gpu_mesh: mesh_instance.gpu_mesh.clone(), + world_from_mesh, + outline_mask_ids, + picking_layer_id: picking_layer_id_from_instance_path_hash( + picking_instance_hash, + ), + ..Default::default() + } + })); - Ok(()) + self.0.add_bounding_box( + entity_path.hash(), + mesh.bbox(), + ent_context.world_from_entity, + ); + }; + } } } @@ -170,18 +155,128 @@ impl VisualizerSystem for Mesh3DVisualizer { fn execute( &mut self, ctx: &ViewerContext<'_>, - query: &ViewQuery<'_>, + view_query: &ViewQuery<'_>, view_ctx: &ViewContextCollection, ) -> Result, SpaceViewSystemExecutionError> { let mut instances = Vec::new(); - process_archetype_views::( + super::entity_iterator::process_archetype::( ctx, - query, + view_query, view_ctx, view_ctx.get::()?.points, - |ctx, ent_path, _ent_props, arch_view, ent_context| { - self.process_arch_view(ctx, &mut instances, &arch_view, ent_path, ent_context) + |ctx, entity_path, _entity_props, spatial_ctx, results| match results { + CachedResults::LatestAt(_query, results) => { + re_tracing::profile_scope!(format!("{entity_path} @ {_query:?}")); + + use crate::visualizers::CachedLatestAtResultsExt as _; + + let resolver = ctx.recording().resolver(); + + let vertex_positions = match results.get_dense::(resolver) { + Some(Ok(positions)) if !positions.is_empty() => positions, + Some(err @ Err(_)) => err?, + _ => return Ok(()), + }; + + let vertex_normals = results.get_or_empty_dense(resolver)?; + let vertex_colors = results.get_or_empty_dense(resolver)?; + let vertex_texcoords = results.get_or_empty_dense(resolver)?; + let mesh_properties = results.get_or_empty_dense(resolver)?; + let mesh_materials = results.get_or_empty_dense(resolver)?; + let albedo_textures = results.get_or_empty_dense(resolver)?; + let class_ids = results.get_or_empty_dense(resolver)?; + + let data = { + // NOTE: + // - Per-vertex properties are joined using the cluster key as usual. + // - Per-mesh properties are just treated as a "global var", essentially. + Mesh3DComponentData { + index: results.compound_index, + + vertex_positions, + vertex_normals, + vertex_colors, + vertex_texcoords, + + mesh_properties: mesh_properties.first(), + mesh_material: mesh_materials.first(), + albedo_texture: albedo_textures.first(), + + class_ids, + } + }; + + self.process_data( + ctx, + &mut instances, + entity_path, + spatial_ctx, + std::iter::once(data), + ); + Ok(()) + } + + CachedResults::Range(_query, results) => { + re_tracing::profile_scope!(format!("{entity_path} @ {_query:?}")); + + use crate::visualizers::CachedRangeResultsExt as _; + + let resolver = ctx.recording().resolver(); + + let vertex_positions = match results.get_dense::(resolver, _query) { + Some(Ok(positions)) => positions, + Some(err @ Err(_)) => err?, + _ => return Ok(()), + }; + + let vertex_normals = results.get_or_empty_dense(resolver, _query)?; + let vertex_colors = results.get_or_empty_dense(resolver, _query)?; + let vertex_texcoords = results.get_or_empty_dense(resolver, _query)?; + let mesh_properties = results.get_or_empty_dense(resolver, _query)?; + let mesh_materials = results.get_or_empty_dense(resolver, _query)?; + let albedo_textures = results.get_or_empty_dense(resolver, _query)?; + let class_ids = results.get_or_empty_dense(resolver, _query)?; + + let data = range_zip_1x7( + vertex_positions.range_indexed(_query.range()), + vertex_normals.range_indexed(_query.range()), + vertex_colors.range_indexed(_query.range()), + vertex_texcoords.range_indexed(_query.range()), + mesh_properties.range_indexed(_query.range()), + mesh_materials.range_indexed(_query.range()), + albedo_textures.range_indexed(_query.range()), + class_ids.range_indexed(_query.range()), + ) + .map( + |( + &index, + vertex_positions, + vertex_normals, + vertex_colors, + vertex_texcoords, + mesh_properties, + mesh_material, + albedo_texture, + class_ids, + )| { + Mesh3DComponentData { + index, + vertex_positions, + vertex_normals: vertex_normals.unwrap_or_default(), + vertex_colors: vertex_colors.unwrap_or_default(), + vertex_texcoords: vertex_texcoords.unwrap_or_default(), + mesh_properties: mesh_properties.and_then(|v| v.first()), + mesh_material: mesh_material.and_then(|v| v.first()), + albedo_texture: albedo_texture.and_then(|v| v.first()), + class_ids: class_ids.unwrap_or_default(), + } + }, + ); + + self.process_data(ctx, &mut instances, entity_path, spatial_ctx, data); + Ok(()) + } }, )?; diff --git a/crates/re_space_view_spatial/src/visualizers/results_ext.rs b/crates/re_space_view_spatial/src/visualizers/results_ext.rs index 3fc880952980..c9338730606b 100644 --- a/crates/re_space_view_spatial/src/visualizers/results_ext.rs +++ b/crates/re_space_view_spatial/src/visualizers/results_ext.rs @@ -4,7 +4,7 @@ use re_query_cache::{ }; use re_types::Component; -// --- +// --- Cached --- pub trait CachedLatestAtResultsExt { fn get_dense<'a, C: Component>( @@ -100,15 +100,11 @@ impl CachedRangeResultsExt for CachedRangeResults { // TODO(#5607): what should happen if the promise is still pending? let (front_status, back_status) = results.status(query.range()); match front_status { - PromiseResult::Error(err) => { - return Err(re_query_cache::QueryError::Other(err.into())) - } + PromiseResult::Error(err) => return Err(re_query_cache::QueryError::Other(err.into())), PromiseResult::Pending | PromiseResult::Ready(_) => {} } match back_status { - PromiseResult::Error(err) => { - return Err(re_query_cache::QueryError::Other(err.into())) - } + PromiseResult::Error(err) => return Err(re_query_cache::QueryError::Other(err.into())), PromiseResult::Pending | PromiseResult::Ready(_) => {} } diff --git a/crates/re_space_view_text_document/src/visualizer_system.rs b/crates/re_space_view_text_document/src/visualizer_system.rs index ea8b404f0a74..75e06692813c 100644 --- a/crates/re_space_view_text_document/src/visualizer_system.rs +++ b/crates/re_space_view_text_document/src/visualizer_system.rs @@ -1,9 +1,6 @@ use re_data_store::LatestAtQuery; -use re_query::{query_archetype, QueryError}; -use re_types::{ - archetypes::{self, TextDocument}, - components, -}; +use re_space_view::external::re_query_cache::PromiseResult; +use re_types::{archetypes::TextDocument, components}; use re_viewer_context::{ IdentifiedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection, ViewQuery, ViewerContext, VisualizerQueryInfo, VisualizerSystem, @@ -37,39 +34,35 @@ impl VisualizerSystem for TextDocumentSystem { fn execute( &mut self, ctx: &ViewerContext<'_>, - query: &ViewQuery<'_>, + view_query: &ViewQuery<'_>, _view_ctx: &ViewContextCollection, ) -> Result, SpaceViewSystemExecutionError> { - let store = ctx.recording_store(); - - let timeline_query = LatestAtQuery::new(query.timeline, query.latest_at); + let timeline_query = LatestAtQuery::new(view_query.timeline, view_query.latest_at); - for data_result in query.iter_visible_data_results(ctx, Self::identifier()) { - // TODO(#3320): this match can go away once the issue is resolved - match query_archetype::( - store, - &timeline_query, - &data_result.entity_path, - ) { - Ok(arch_view) => { - let bodies = arch_view.iter_required_component::()?; - let media_types = - arch_view.iter_optional_component::()?; - - for (body, media_type) in itertools::izip!(bodies, media_types) { - let media_type = media_type.unwrap_or(components::MediaType::plain_text()); - self.text_entries - .push(TextDocumentEntry { body, media_type }); - } + for data_result in view_query.iter_visible_data_results(ctx, Self::identifier()) { + let TextDocument { text, media_type } = match ctx + .recording() + .latest_at_archetype(&data_result.entity_path, &timeline_query) + { + PromiseResult::Pending | PromiseResult::Ready(None) => { + // TODO(#5607): what should happen if the promise is still pending? + continue; } - Err(QueryError::PrimaryNotFound(_)) => {} - Err(err) => { + PromiseResult::Ready(Some((_, arch))) => arch, + PromiseResult::Error(err) => { re_log::error_once!( "Unexpected error querying {:?}: {err}", &data_result.entity_path ); + continue; } }; + + let media_type = media_type.unwrap_or(components::MediaType::plain_text()); + self.text_entries.push(TextDocumentEntry { + body: text, + media_type, + }); } Ok(Vec::new()) diff --git a/crates/re_viewer/src/lib.rs b/crates/re_viewer/src/lib.rs index 1c9b094191ad..85b2747bf3e1 100644 --- a/crates/re_viewer/src/lib.rs +++ b/crates/re_viewer/src/lib.rs @@ -34,8 +34,8 @@ pub mod external { pub use {eframe, egui}; pub use { re_data_store, re_data_store::external::*, re_entity_db, re_log, re_log_types, re_memory, - re_renderer, re_types, re_ui, re_viewer_context, re_viewer_context::external::*, - re_viewport, re_viewport::external::*, + re_query_cache, re_renderer, re_types, re_ui, re_viewer_context, + re_viewer_context::external::*, re_viewport, re_viewport::external::*, }; } diff --git a/crates/re_viewer_context/Cargo.toml b/crates/re_viewer_context/Cargo.toml index a4bae350f3e4..f8273dda4419 100644 --- a/crates/re_viewer_context/Cargo.toml +++ b/crates/re_viewer_context/Cargo.toml @@ -23,7 +23,6 @@ re_log_types.workspace = true re_log.workspace = true re_query_cache.workspace = true re_query.workspace = true -re_query2.workspace = true re_renderer.workspace = true re_smart_channel.workspace = true re_string_interner.workspace = true diff --git a/crates/re_viewer_context/src/annotations.rs b/crates/re_viewer_context/src/annotations.rs index 1853690a5e0d..f51010ffb89b 100644 --- a/crates/re_viewer_context/src/annotations.rs +++ b/crates/re_viewer_context/src/annotations.rs @@ -6,7 +6,6 @@ use nohash_hasher::IntSet; use re_data_store::LatestAtQuery; use re_entity_db::EntityPath; use re_log_types::RowId; -use re_query::{query_archetype, ArchetypeView}; use re_types::archetypes::AnnotationContext; use re_types::datatypes::{AnnotationInfo, ClassDescription, ClassId, KeypointId, Utf8}; @@ -38,29 +37,6 @@ impl Annotations { .clone() } - pub fn try_from_view(view: &ArchetypeView) -> Option { - re_tracing::profile_function!(); - - use re_log::ResultExt as _; - - view.optional_mono_component::() - .warn_on_err_once("Failed to load AnnotationContext") - .flatten() - .map(|ctx| Self { - row_id: view.primary_row_id(), - class_map: ctx - .0 - .into_iter() - .map(|elem| { - ( - elem.class_id, - CachedClassDescription::from(elem.class_description), - ) - }) - .collect(), - }) - } - #[inline] pub fn resolved_class_description( &self, @@ -253,8 +229,6 @@ impl AnnotationMap { let mut visited = IntSet::::default(); - let data_store = ctx.recording_store(); - // This logic is borrowed from `iter_ancestor_meta_field`, but using the arrow-store instead // not made generic as `AnnotationContext` was the only user of that function for ent_path in entities { @@ -273,13 +247,27 @@ impl AnnotationMap { // Otherwise check the obj_store for the field. // If we find one, insert it and then we can break. std::collections::btree_map::Entry::Vacant(entry) => { - if query_archetype::(data_store, time_query, &parent) + if let Some(((_, row_id), ann_ctx)) = ctx + .recording() + .latest_at_archetype::(&parent, time_query) .ok() - .and_then(|view| Annotations::try_from_view(&view)) - .map(|annotations| entry.insert(Arc::new(annotations))) - .is_some() + .flatten() { - break; + let annotations = Annotations { + row_id, + class_map: ann_ctx + .context + .0 + .into_iter() + .map(|elem| { + ( + elem.class_id, + CachedClassDescription::from(elem.class_description), + ) + }) + .collect(), + }; + entry.insert(Arc::new(annotations)); } } } diff --git a/crates/re_viewer_context/src/item.rs b/crates/re_viewer_context/src/item.rs index e466b0e7e152..0d21d88d3dfd 100644 --- a/crates/re_viewer_context/src/item.rs +++ b/crates/re_viewer_context/src/item.rs @@ -1,4 +1,4 @@ -use re_entity_db::InstancePath; +use re_entity_db::{EntityDb, InstancePath}; use re_log_types::{ComponentPath, DataPath, EntityPath}; use crate::{ContainerId, SpaceViewId}; @@ -151,18 +151,18 @@ impl Item { /// If the given item refers to the first element of an instance with a single element, resolve to a splatted entity path. pub fn resolve_mono_instance_path_item( + entity_db: &EntityDb, query: &re_data_store::LatestAtQuery, - store: &re_data_store::DataStore, item: &Item, ) -> Item { // Resolve to entity path if there's only a single instance. match item { Item::InstancePath(instance_path) => { - Item::InstancePath(resolve_mono_instance_path(query, store, instance_path)) + Item::InstancePath(resolve_mono_instance_path(entity_db, query, instance_path)) } Item::DataResult(space_view_id, instance_path) => Item::DataResult( *space_view_id, - resolve_mono_instance_path(query, store, instance_path), + resolve_mono_instance_path(entity_db, query, instance_path), ), Item::AppId(_) | Item::DataSource(_) @@ -175,27 +175,36 @@ pub fn resolve_mono_instance_path_item( /// If the given path refers to the first element of an instance with a single element, resolve to a splatted entity path. pub fn resolve_mono_instance_path( + entity_db: &EntityDb, query: &re_data_store::LatestAtQuery, - store: &re_data_store::DataStore, instance: &re_entity_db::InstancePath, ) -> re_entity_db::InstancePath { re_tracing::profile_function!(); if instance.instance_key.0 == 0 { - let Some(components) = store.all_components(&query.timeline(), &instance.entity_path) + // NOTE: While we normally frown upon direct queries to the datastore, `all_components` is fine. + let Some(components) = entity_db + .store() + .all_components(&query.timeline(), &instance.entity_path) else { // No components at all, return splatted entity. return re_entity_db::InstancePath::entity_splat(instance.entity_path.clone()); }; + for component in components { - if let Some((_, _row_id, instances)) = re_query::get_component_with_instances( - store, + let results = entity_db.query_caches().latest_at( + entity_db.store(), query, &instance.entity_path, - component, - ) { - if instances.len() > 1 { - return instance.clone(); + [component], + ); + if let Some(results) = results.get(component) { + if let re_query_cache::PromiseResult::Ready(cell) = + results.resolved(entity_db.resolver()) + { + if cell.num_instances() > 1 { + return instance.clone(); + } } } } diff --git a/crates/re_viewer_context/src/selection_state.rs b/crates/re_viewer_context/src/selection_state.rs index 8a57b8aa0421..54ab55c2c337 100644 --- a/crates/re_viewer_context/src/selection_state.rs +++ b/crates/re_viewer_context/src/selection_state.rs @@ -114,8 +114,8 @@ impl ItemCollection { .map(|(item, space_ctx)| { ( resolve_mono_instance_path_item( + ctx.recording(), &ctx.current_query(), - ctx.recording_store(), &item, ), space_ctx, diff --git a/crates/re_viewer_context/src/space_view/mod.rs b/crates/re_viewer_context/src/space_view/mod.rs index 61c60dfadf77..17632a2aa4ff 100644 --- a/crates/re_viewer_context/src/space_view/mod.rs +++ b/crates/re_viewer_context/src/space_view/mod.rs @@ -49,7 +49,7 @@ pub enum SpaceViewSystemExecutionError { QueryError(#[from] re_query::QueryError), #[error(transparent)] - QueryError2(#[from] re_query2::QueryError), + QueryError2(#[from] re_query_cache::QueryError), #[error(transparent)] DeserializationError(#[from] re_types::DeserializationError), diff --git a/crates/re_viewport/src/container.rs b/crates/re_viewport/src/container.rs index 65f53e480bd8..9e663aceeb29 100644 --- a/crates/re_viewport/src/container.rs +++ b/crates/re_viewport/src/container.rs @@ -2,7 +2,7 @@ use ahash::HashMap; use egui_tiles::TileId; use re_data_store::LatestAtQuery; -use re_entity_db::{external::re_query2::PromiseResult, EntityDb}; +use re_entity_db::{external::re_query_cache::PromiseResult, EntityDb}; use re_log::ResultExt; use re_log_types::{DataRow, EntityPath, RowId}; use re_types::blueprint::components::Visible; @@ -59,7 +59,7 @@ impl ContainerBlueprint { // TODO(#5607): what should happen if the promise is still pending? None } - PromiseResult::Ready(arch) => arch, + PromiseResult::Ready(arch) => arch.map(|(_, arch)| arch), PromiseResult::Error(err) => { if cfg!(debug_assertions) { re_log::error!("Failed to load container blueprint: {err}."); diff --git a/crates/re_viewport/src/viewport_blueprint.rs b/crates/re_viewport/src/viewport_blueprint.rs index 241092034d29..f0cb5f83d587 100644 --- a/crates/re_viewport/src/viewport_blueprint.rs +++ b/crates/re_viewport/src/viewport_blueprint.rs @@ -4,10 +4,10 @@ use std::sync::atomic::{AtomicBool, Ordering}; use ahash::HashMap; use egui_tiles::{SimplificationOptions, TileId}; use nohash_hasher::IntSet; -use re_entity_db::external::re_query2::PromiseResult; use smallvec::SmallVec; use re_data_store::LatestAtQuery; +use re_entity_db::external::re_query_cache::PromiseResult; use re_entity_db::EntityPath; use re_space_view::SpaceViewBlueprint; use re_types::blueprint::components::ViewerRecommendationHash; @@ -80,7 +80,7 @@ impl ViewportBlueprint { // TODO(#5607): what should happen if the promise is still pending? Default::default() } - PromiseResult::Ready(arch) => arch.unwrap_or_default(), + PromiseResult::Ready(arch) => arch.map_or_else(Default::default, |(_, arch)| arch), PromiseResult::Error(err) => { if cfg!(debug_assertions) { re_log::error!("Failed to load viewport blueprint: {err}."); diff --git a/examples/rust/custom_space_view/Cargo.toml b/examples/rust/custom_space_view/Cargo.toml index 952646cfa948..e7f0bd820a78 100644 --- a/examples/rust/custom_space_view/Cargo.toml +++ b/examples/rust/custom_space_view/Cargo.toml @@ -14,6 +14,7 @@ analytics = ["re_crash_handler/analytics", "re_viewer/analytics"] [dependencies] re_crash_handler = { path = "../../../crates/re_crash_handler" } +re_query2 = { path = "../../../crates/re_query2" } re_viewer = { path = "../../../crates/re_viewer", default-features = false } # We need re_sdk_comms to receive log events from and SDK: diff --git a/examples/rust/custom_space_view/src/color_coordinates_visualizer_system.rs b/examples/rust/custom_space_view/src/color_coordinates_visualizer_system.rs index 66a37ce5c786..a90e12136fb1 100644 --- a/examples/rust/custom_space_view/src/color_coordinates_visualizer_system.rs +++ b/examples/rust/custom_space_view/src/color_coordinates_visualizer_system.rs @@ -1,9 +1,12 @@ use re_viewer::external::{ egui, re_log_types::EntityPath, - re_query::query_archetype, - re_renderer, - re_types::{self, components::InstanceKey, ComponentName, Loggable as _}, + re_query_cache, re_renderer, + re_types::{ + self, + components::{Color, InstanceKey}, + ComponentName, Loggable as _, + }, re_viewer_context::{ IdentifiedViewSystem, SpaceViewSystemExecutionError, ViewContextCollection, ViewQuery, ViewSystemIdentifier, ViewerContext, VisualizerQueryInfo, VisualizerSystem, @@ -35,6 +38,16 @@ impl re_types::Archetype for ColorArchetype { } } +impl re_query2::ToArchetype for re_query_cache::CachedLatestAtResults { + #[inline] + fn to_archetype( + &self, + _resolver: &re_query2::PromiseResolver, + ) -> re_query2::PromiseResult> { + re_query2::PromiseResult::Ready(Ok(ColorArchetype)) + } +} + impl IdentifiedViewSystem for InstanceColorSystem { fn identifier() -> ViewSystemIdentifier { "InstanceColor".into() @@ -56,30 +69,36 @@ impl VisualizerSystem for InstanceColorSystem { // For each entity in the space view that should be displayed with the `InstanceColorSystem`… for data_result in query.iter_visible_data_results(ctx, Self::identifier()) { // …gather all colors and their instance ids. - if let Ok(arch_view) = query_archetype::( + + let results = re_query2::latest_at( ctx.recording_store(), &ctx.current_query(), &data_result.entity_path, - ) { - if let Ok(colors) = - arch_view.iter_required_component::() - { - self.colors.push(( - data_result.entity_path.clone(), - arch_view - .iter_instance_keys() - .zip(colors) - .map(|(instance_key, color)| { - let [r, g, b, _] = color.to_array(); - ColorWithInstanceKey { - color: egui::Color32::from_rgb(r, g, b), - instance_key, - } - }) - .collect(), - )); - } - } + [Color::name()], + ); + + let Some(colors) = results.get(Color::name()).and_then(|results| { + results + .to_dense::(ctx.recording().resolver()) + .flatten() + .ok() + }) else { + continue; + }; + + self.colors.push(( + data_result.entity_path.clone(), + (0..) + .zip(colors) + .map(|(instance_key, color)| { + let [r, g, b, _] = color.to_array(); + ColorWithInstanceKey { + color: egui::Color32::from_rgb(r, g, b), + instance_key: instance_key.into(), + } + }) + .collect(), + )); } // We're not using `re_renderer` here, so return an empty vector.