Implement sizebytes for things in the tf transform cache (#12038)

Author: IsseW

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(),
         }