Implement sizebytes for things in the tf transform cache

Cargo.lock

@@ -8520,6 +8520,7 @@ name = "re_byte_size"
 version = "0.28.0-alpha.1+dev"
 dependencies = [
  "arrow",
+ "glam",
  "half",
  "smallvec",
 ]
@@ -9745,6 +9746,7 @@ dependencies = [
  "nohash-hasher",
  "parking_lot",
  "re_arrow_util",
+ "re_byte_size",
  "re_chunk_store",
  "re_entity_db",
  "re_log",

crates/store/re_tf/Cargo.toml

@@ -21,6 +21,7 @@ all-features = true
 
 [dependencies]
 re_arrow_util.workspace = true
+re_byte_size = { workspace = true, features = ["glam"] }
 re_chunk_store.workspace = true
 re_entity_db.workspace = true # It would be nice not to depend on this, but we need this in order to do queries right now.
 re_log.workspace = true

crates/store/re_tf/src/frame_id_registry.rs

@@ -1,9 +1,9 @@
 use std::collections::hash_map::Entry;
 
 use nohash_hasher::IntMap;
+use re_byte_size::SizeBytes;
 use re_log_types::EntityPath;
-use re_types::components::TransformFrameId;
-use re_types::{TransformFrameIdHash, archetypes};
+use re_types::{TransformFrameIdHash, archetypes, components::TransformFrameId};
 
 /// Frame id registry for resolving frame id hashes back to frame ids.
 pub struct FrameIdRegistry {
@@ -25,6 +25,16 @@ impl Default for FrameIdRegistry {
  }
 }
 
+impl SizeBytes for FrameIdRegistry {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ frame_id_lookup_table,
+ } = self;
+
+ frame_id_lookup_table.total_size_bytes()
+ }
+}
+
 impl FrameIdRegistry {
  /// Looks up a frame ID by its hash.
  ///

crates/store/re_tf/src/transform_resolution_cache.rs

@@ -5,6 +5,7 @@ use glam::DAffine3;
 use itertools::{Either, izip};
 use nohash_hasher::IntMap;
 use parking_lot::Mutex;
+use re_byte_size::SizeBytes;
 
 use crate::frame_id_registry::FrameIdRegistry;
 use crate::{
@@ -17,8 +18,7 @@ use crate::{
 };
 
 use re_arrow_util::ArrowArrayDowncastRef as _;
-use re_chunk_store::external::arrow;
-use re_chunk_store::{Chunk, LatestAtQuery};
+use re_chunk_store::{Chunk, LatestAtQuery, external::arrow};
 use re_entity_db::EntityDb;
 use re_log_types::{EntityPath, TimeInt, TimelineName};
 use re_types::{ComponentIdentifier, archetypes, components};
@@ -61,6 +61,20 @@ impl Default for TransformResolutionCache {
  }
 }
 
+impl SizeBytes for TransformResolutionCache {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ frame_id_registry,
+ per_timeline,
+ static_timeline,
+ } = self;
+
+ frame_id_registry.heap_size_bytes()
+ + per_timeline.heap_size_bytes()
+ + static_timeline.heap_size_bytes()
+ }
+}
+
 /// A transform from a child frame to a parent frame.
 #[derive(Clone, Debug, PartialEq)]
 pub struct ParentFromChildTransform {
@@ -71,6 +85,14 @@ pub struct ParentFromChildTransform {
  pub transform: DAffine3,
 }
 
+impl SizeBytes for ParentFromChildTransform {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self { parent, transform } = self;
+
+ parent.heap_size_bytes() + transform.heap_size_bytes()
+ }
+}
+
 /// Cached transforms for a single timeline.
 ///
 /// Includes any static transforms that may apply globally.
@@ -243,6 +265,20 @@ impl CachedTransformsForTimeline {
  }
 }
 
+impl SizeBytes for CachedTransformsForTimeline {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ per_child_frame_transforms,
+ non_recursive_clears,
+ recursive_clears,
+ } = self;
+
+ per_child_frame_transforms.heap_size_bytes()
+ + non_recursive_clears.heap_size_bytes()
+ + recursive_clears.heap_size_bytes()
+ }
+}
+
 #[derive(Clone, Debug, PartialEq)]
 enum CachedTransformValue<T> {
  /// Cache is invalidated, we don't know what state we're in.
@@ -255,6 +291,15 @@ enum CachedTransformValue<T> {
  Cleared,
 }
 
+impl<T: SizeBytes> SizeBytes for CachedTransformValue<T> {
+ fn heap_size_bytes(&self) -> u64 {
+ match self {
+ Self::Resident(item) => item.heap_size_bytes(),
+ Self::Invalidated | Self::Cleared => 0,
+ }
+ }
+}
+
 type FrameTransformTimeMap = BTreeMap<TimeInt, CachedTransformValue<ParentFromChildTransform>>;
 
 type PoseTransformTimeMap = BTreeMap<TimeInt, CachedTransformValue<Vec<DAffine3>>>;
@@ -344,6 +389,20 @@ impl TransformsForChildFrameEvents {
  }
 }
 
+impl SizeBytes for TransformsForChildFrameEvents {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ frame_transforms,
+ pose_transforms,
+ pinhole_projections,
+ } = self;
+
+ frame_transforms.heap_size_bytes()
+ + pose_transforms.heap_size_bytes()
+ + pinhole_projections.heap_size_bytes()
+ }
+}
+
 /// Cached transforms from a single child frame to a (potentially changing) parent frame over time.
 ///
 /// Incorporates any static transforms that may apply to this entity.
@@ -391,6 +450,23 @@ impl PartialEq for TransformsForChildFrame {
  }
 }
 
+impl SizeBytes for TransformsForChildFrame {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ associated_entity_path,
+ child_frame,
+ events,
+
+ #[cfg(debug_assertions)]
+ timeline: _,
+ } = self;
+
+ associated_entity_path.heap_size_bytes()
+ + child_frame.heap_size_bytes()
+ + events.lock().heap_size_bytes()
+ }
+}
+
 fn add_invalidated_entry_if_not_already_cleared<T: PartialEq>(
  transforms: &mut BTreeMap<TimeInt, CachedTransformValue<T>>,
  time: TimeInt,
@@ -424,6 +500,22 @@ pub struct ResolvedPinholeProjection {
  pub view_coordinates: components::ViewCoordinates,
 }
 
+impl SizeBytes for ResolvedPinholeProjection {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ parent,
+ image_from_camera,
+ resolution,
+ view_coordinates,
+ } = self;
+
+ parent.heap_size_bytes()
+ + image_from_camera.heap_size_bytes()
+ + resolution.heap_size_bytes()
+ + view_coordinates.heap_size_bytes()
+ }
+}
+
 impl TransformsForChildFrame {
  fn new(
  associated_entity_path: EntityPath,

crates/store/re_types/Cargo.toml

@@ -34,7 +34,7 @@ ecolor = ["dep:ecolor"]
 egui_plot = ["dep:egui_plot"]
 
 ## Add support for some math operations using [`glam`](https://crates.io/crates/glam/) and [`macaw`](https://crates.io/crates/macaw/).
-glam = ["dep:glam", "dep:macaw"]
+glam = ["dep:glam", "dep:macaw", "re_byte_size/glam"]
 
 ## Integration with the [`image`](https://crates.io/crates/image/) crate, plus JPEG and TIFF support.
 image = ["dep:ecolor", "dep:image", "dep:tiff"]

crates/top/rerun/Cargo.toml

@@ -66,7 +66,7 @@ ecolor = ["re_types?/ecolor"]
 
 ## Add support for some math operations using [`glam`](https://crates.io/crates/glam/).
 ## Only relevant if feature `sdk` is enabled.
-glam = ["re_types?/glam"]
+glam = ["re_types?/glam", "re_byte_size/glam"]
 
 ## Integration with the [`image`](https://crates.io/crates/image/) crate, plus JPEG support.
 image = ["re_types?/image"]

crates/utils/re_byte_size/Cargo.toml

@@ -20,10 +20,12 @@ all-features = true
 
 
 [features]
+glam = ["dep:glam"]
 
 
 [dependencies]
 # TODO(emilk): make some of these opt-in
 arrow.workspace = true
 half.workspace = true
 smallvec.workspace = true
+glam = { workspace = true, optional = true }

crates/utils/re_byte_size/src/primitive_sizes.rs

@@ -28,3 +28,6 @@ impl_size_bytes_pod!(
  u8, u16, u32, u64, u128, i8, i16, i32, i64, i128, bool, f32, f64
 );
 impl_size_bytes_pod!(half::f16);
+
+#[cfg(feature = "glam")]
+impl_size_bytes_pod!(glam::DAffine3);

crates/viewer/re_viewer_context/src/cache/transform_database_store.rs

@@ -1,4 +1,5 @@
 use parking_lot::{ArcMutexGuard, Mutex, RawMutex};
+use re_byte_size::SizeBytes;
 use std::sync::Arc;
 
 use super::{Cache, CacheMemoryReport};
@@ -34,6 +35,18 @@ impl TransformDatabaseStoreCache {
  }
 }
 
+impl SizeBytes for TransformDatabaseStoreCache {
+ fn heap_size_bytes(&self) -> u64 {
+ let Self {
+ initialized,
+ transform_cache,
+ } = self;
+ let cache = transform_cache.lock();
+
+ initialized.heap_size_bytes() + cache.total_size_bytes()
+ }
+}
+
 impl Cache for TransformDatabaseStoreCache {
  fn purge_memory(&mut self) {
  // Can't purge memory from the transform cache right now and even if we could, there's
@@ -42,8 +55,7 @@ impl Cache for TransformDatabaseStoreCache {
 
  fn memory_report(&self) -> CacheMemoryReport {
  CacheMemoryReport {
- // TODO(RR-2517): Implement SizeBytes for TransformResolutionCache.
- bytes_cpu: 0, //self.transform_cache.total_size_bytes(),
+ bytes_cpu: self.total_size_bytes(),
  bytes_gpu: None,
  per_cache_item_info: Vec::new(),
  }