Implement sizebytes for things in the tf transform cache

Cargo.lock

@@ -8518,6 +8518,7 @@ name = "re_byte_size"
 version = "0.28.0-alpha.1+dev"
 dependencies = [
  "arrow",
+ "glam",
  "half",
  "smallvec",
 ]

crates/store/re_tf/src/frame_id_registry.rs

@@ -1,6 +1,7 @@
 use std::collections::hash_map::Entry;
 
 use nohash_hasher::IntMap;
+use re_chunk_store::external::re_chunk::external::re_byte_size::SizeBytes;
 use re_log_types::EntityPath;
 use re_types::components::TransformFrameId;
 use re_types::{TransformFrameIdHash, archetypes};
@@ -25,6 +26,12 @@ impl Default for FrameIdRegistry {
     }
 }
 
+impl SizeBytes for FrameIdRegistry {
+    fn heap_size_bytes(&self) -> u64 {
+        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_chunk_store::external::re_chunk::external::re_byte_size::SizeBytes;
 
 use crate::frame_id_registry::FrameIdRegistry;
 use crate::{
@@ -61,6 +62,14 @@ impl Default for TransformResolutionCache {
     }
 }
 
+impl SizeBytes for TransformResolutionCache {
+    fn heap_size_bytes(&self) -> u64 {
+        self.frame_id_registry.heap_size_bytes()
+            + self.per_timeline.heap_size_bytes()
+            + self.static_timeline.heap_size_bytes()
+    }
+}
+
 /// A transform from a child frame to a parent frame.
 #[derive(Clone, Debug, PartialEq)]
 pub struct ParentFromChildTransform {
@@ -71,6 +80,12 @@ pub struct ParentFromChildTransform {
     pub transform: DAffine3,
 }
 
+impl SizeBytes for ParentFromChildTransform {
+    fn heap_size_bytes(&self) -> u64 {
+        self.parent.heap_size_bytes() + self.transform.heap_size_bytes()
+    }
+}
+
 /// Cached transforms for a single timeline.
 ///
 /// Includes any static transforms that may apply globally.
@@ -243,6 +258,14 @@ impl CachedTransformsForTimeline {
     }
 }
 
+impl SizeBytes for CachedTransformsForTimeline {
+    fn heap_size_bytes(&self) -> u64 {
+        self.per_child_frame_transforms.heap_size_bytes()
+            + self.non_recursive_clears.heap_size_bytes()
+            + self.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 +278,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 +376,14 @@ impl TransformsForChildFrameEvents {
     }
 }
 
+impl SizeBytes for TransformsForChildFrameEvents {
+    fn heap_size_bytes(&self) -> u64 {
+        self.frame_transforms.heap_size_bytes()
+            + self.pose_transforms.heap_size_bytes()
+            + self.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 +431,15 @@ impl PartialEq for TransformsForChildFrame {
     }
 }
 
+impl SizeBytes for TransformsForChildFrame {
+    fn heap_size_bytes(&self) -> u64 {
+        self.timeline.heap_size_bytes()
+            + self.associated_entity_path.heap_size_bytes()
+            + self.child_frame.heap_size_bytes()
+            + self.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 +473,15 @@ pub struct ResolvedPinholeProjection {
     pub view_coordinates: components::ViewCoordinates,
 }
 
+impl SizeBytes for ResolvedPinholeProjection {
+    fn heap_size_bytes(&self) -> u64 {
+        self.parent.heap_size_bytes()
+            + self.image_from_camera.heap_size_bytes()
+            + self.resolution.heap_size_bytes()
+            + self.view_coordinates.heap_size_bytes()
+    }
+}
+
 impl TransformsForChildFrame {
     fn new(
         associated_entity_path: EntityPath,

crates/utils/re_byte_size/Cargo.toml

@@ -27,3 +27,4 @@ all-features = true
 arrow.workspace = true
 half.workspace = true
 smallvec.workspace = true
+glam.workspace = true

crates/utils/re_byte_size/src/primitive_sizes.rs

@@ -28,3 +28,4 @@ impl_size_bytes_pod!(
     u8, u16, u32, u64, u128, i8, i16, i32, i64, i128, bool, f32, f64
 );
 impl_size_bytes_pod!(half::f16);
+impl_size_bytes_pod!(glam::DAffine3);

crates/viewer/re_view_spatial/src/caches/transform_database_store.rs

@@ -1,4 +1,5 @@
 use parking_lot::{ArcMutexGuard, Mutex, RawMutex};
+use re_byte_size::SizeBytes;
 use std::sync::Arc;
 
 use re_chunk_store::ChunkStoreEvent;
@@ -34,6 +35,14 @@ impl TransformDatabaseStoreCache {
     }
 }
 
+impl SizeBytes for TransformDatabaseStoreCache {
+    fn heap_size_bytes(&self) -> u64 {
+        let cache = self.transform_cache.lock();
+
+        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 +51,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(),
         }