Introduce DataSizeEstimator trait in preparation for supporting bevy builtins.

Author: BGR360

src/app_ext.rs

@@ -3,10 +3,10 @@ use std::any::Any;
 use bevy::{
     app::App,
     asset::Asset,
-    ecs::{component::Component, system::Resource},
+    ecs::{component::Component, schedule::IntoSystemDescriptor, system::Resource},
 };
 
-use crate::{systems, DataSize, MemoryUsage};
+use crate::{systems, DataSize, DataSizeEstimator, MemoryUsage};
 
 /// [`App`] extension methods to register types for memory usage tracking.
 pub trait RegisterSizedTypes {
@@ -17,7 +17,10 @@ pub trait RegisterSizedTypes {
     /// resource.
     fn register_sized_component<T>(&mut self) -> &mut Self
     where
-        T: Any + DataSize + Component;
+        T: Any + DataSize + Component,
+    {
+        self.register_sized_type::<T, T, _, _>(systems::update_stats_for_component::<T, T>)
+    }
 
     /// Registers the given [`Resource`] type with the
     /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
@@ -26,7 +29,10 @@ pub trait RegisterSizedTypes {
     /// resource.
     fn register_sized_resource<T>(&mut self) -> &mut Self
     where
-        T: Any + DataSize + Resource;
+        T: Any + DataSize + Resource,
+    {
+        self.register_sized_type::<T, T, _, _>(systems::update_stats_for_resource::<T, T>)
+    }
 
     /// Registers the given [`Asset`] type with the
     /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
@@ -35,47 +41,38 @@ pub trait RegisterSizedTypes {
     /// resource.
     fn register_sized_asset<T>(&mut self) -> &mut Self
     where
-        T: Any + DataSize + Asset;
-}
-
-impl RegisterSizedTypes for App {
-    fn register_sized_component<T>(&mut self) -> &mut Self
-    where
-        T: Any + Component + DataSize,
+        T: Any + DataSize + Asset,
     {
-        register_type::<T>(self);
-
-        self.add_system(systems::update_stats_for_component::<T>);
-
-        self
+        self.register_sized_type::<T, T, _, _>(systems::update_stats_for_asset::<T, T>)
     }
 
-    fn register_sized_resource<T>(&mut self) -> &mut Self
+    /// Registers a type that can be sized with the given [`DataSizeEstimator`]
+    /// type and updated using the provided `system`.
+    fn register_sized_type<T, E, S, Params>(&mut self, system: S) -> &mut Self
     where
-        T: Any + DataSize + Resource,
-    {
-        register_type::<T>(self);
-
-        self.add_system(systems::update_stats_for_resource::<T>);
-
-        self
-    }
+        T: Any,
+        E: DataSizeEstimator<T>,
+        S: IntoSystemDescriptor<Params>;
+}
 
-    fn register_sized_asset<T>(&mut self) -> &mut Self
+impl RegisterSizedTypes for App {
+    fn register_sized_type<T, E, S, Params>(&mut self, system: S) -> &mut Self
     where
-        T: Any + DataSize + Asset,
+        T: Any,
+        E: DataSizeEstimator<T>,
+        S: IntoSystemDescriptor<Params>,
     {
         register_type::<T>(self);
 
-        self.add_system(systems::update_stats_for_asset::<T>);
+        self.add_system(system);
 
         self
     }
 }
 
 fn register_type<T>(app: &mut App)
 where
-    T: Any + DataSize,
+    T: Any,
 {
     let mut memory_usage = app.world.get_resource_mut::<MemoryUsage>().expect(
         "Cannot find resource `MemoryUsage`. Did you forget to add the `MemoryUsagePlugin`?",

src/estimator.rs

@@ -0,0 +1,22 @@
+use crate::DataSize;
+
+/// Indicates that a type can estimate the heap usage of values of type `T`.
+///
+/// All types which implement [`DataSize`] have this trait implemented on
+/// themselves.
+///
+/// This trait exists in order to work around the fact that this crate cannot
+/// add impls of [`DataSize`] to types from the Bevy crate(s).
+pub trait DataSizeEstimator<T> {
+    /// Estimates the size of heap memory taken up by the given value.
+    ///
+    /// Does not include data on the stack, which is usually determined using
+    /// [`std::mem::size_of`].
+    fn estimate_heap_size(value: &T) -> usize;
+}
+
+impl<T: DataSize> DataSizeEstimator<T> for T {
+    fn estimate_heap_size(value: &T) -> usize {
+        <T as DataSize>::estimate_heap_size(value)
+    }
+}

src/lib.rs

@@ -44,14 +44,17 @@ pub use datasize;
 pub use datasize::DataSize;
 
 mod app_ext;
+pub mod builtins;
 mod config;
+mod estimator;
 mod plugin;
 mod resource;
 mod stats;
 pub mod systems;
 
 pub use app_ext::RegisterSizedTypes;
 pub use config::MemoryConfig;
+pub use estimator::DataSizeEstimator;
 pub use plugin::MemoryUsagePlugin;
 pub use resource::MemoryUsage;
 pub use stats::MemoryStats;

src/stats.rs

@@ -3,7 +3,7 @@ use std::{
     sync::atomic::{AtomicUsize, Ordering},
 };
 
-use crate::DataSize;
+use crate::{DataSize, DataSizeEstimator};
 
 /// Memory usage statistics for a single data type.
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
@@ -39,11 +39,22 @@ impl MemoryStats {
     pub fn from_value<T>(value: &T) -> Self
     where
         T: Any + DataSize,
+    {
+        Self::from_value_with_estimator::<T, T>(value)
+    }
+
+    /// Returns the computed memory statistics for a single value using a
+    /// specific [`DataSizeEstimator`].
+    #[inline]
+    pub fn from_value_with_estimator<T, E>(value: &T) -> Self
+    where
+        T: Any,
+        E: DataSizeEstimator<T>,
     {
         Self {
             count: 1,
             total_stack_bytes: Self::stack_size_of(value),
-            total_heap_bytes: Self::heap_size_of(value),
+            total_heap_bytes: Self::heap_size_of_with_estimator::<T, E>(value),
         }
     }
 
@@ -53,6 +64,17 @@ impl MemoryStats {
     where
         T: Any + DataSize,
         I: IntoIterator<Item = &'a T>,
+    {
+        Self::from_values_with_estimator::<T, T, I>(values)
+    }
+
+    /// Returns the computed memory statistics for a collection of values.
+    #[inline]
+    pub fn from_values_with_estimator<'a, T, E, I>(values: I) -> Self
+    where
+        T: Any,
+        E: DataSizeEstimator<T>,
+        I: IntoIterator<Item = &'a T>,
     {
         let mut count = 0;
         let mut total_stack_bytes = 0;
@@ -61,7 +83,7 @@ impl MemoryStats {
         for value in values.into_iter() {
             count += 1;
             total_stack_bytes += Self::stack_size_of(value);
-            total_heap_bytes += Self::heap_size_of(value);
+            total_heap_bytes += Self::heap_size_of_with_estimator::<T, E>(value);
         }
 
         Self {
@@ -94,20 +116,17 @@ impl MemoryStats {
     where
         T: DataSize,
     {
-        value.estimate_heap_size()
+        Self::heap_size_of_with_estimator::<T, T>(value)
     }
 
-    /// Returns the estimated total size of the given value.
-    ///
-    /// This quantity is the sum of [`stack_size_of`] and [`heap_size_of`].
-    ///
-    /// [`stack_size_of`]: Self::stack_size_of
-    /// [`heap_size_of`]: Self::heap_size_of
-    pub fn total_size_of<T>(value: &T) -> usize
+    /// Returns the estimated heap size of the given value using a specific
+    /// [`DataSizeEstimator`].
+    #[inline]
+    pub fn heap_size_of_with_estimator<T, E>(value: &T) -> usize
     where
-        T: Any + DataSize,
+        E: DataSizeEstimator<T>,
     {
-        Self::stack_size_of(value) + Self::heap_size_of(value)
+        <E as DataSizeEstimator<T>>::estimate_heap_size(value)
     }
 }
 

src/systems.rs

@@ -9,53 +9,59 @@ use bevy::{
         system::{Query, Res, Resource},
     },
 };
-use datasize::DataSize;
 
-use crate::{MemoryConfig, MemoryStats, MemoryUsage};
+use crate::{DataSizeEstimator, MemoryConfig, MemoryStats, MemoryUsage};
 
-/// This system iterates through all components of type `T` and updates their
-/// [`MemoryStats`].
-pub fn update_stats_for_component<T>(
+/// This system updates the [`MemoryStats`] for the given component type `T`
+/// using the given [`DataSizeEstimator`].
+pub fn update_stats_for_component<T, E>(
     query: Query<&T>,
     memory_config: Res<MemoryConfig>,
     memory_usage: Res<MemoryUsage>,
 ) where
-    T: Any + Component + DataSize,
+    T: Any + Component,
+    E: DataSizeEstimator<T>,
 {
     update_stats::<T, _>(memory_config, memory_usage, || {
-        MemoryStats::from_values(query.iter())
+        MemoryStats::from_values_with_estimator::<T, E, _>(query.iter())
     });
 }
 
-/// This system updates the [`MemoryStats`] for the given resource type `T`.
-pub fn update_stats_for_resource<T>(
+/// This system updates the [`MemoryStats`] for the given resource type `T`
+/// using the given [`DataSizeEstimator`].
+pub fn update_stats_for_resource<T, E>(
     resource: Res<T>,
     memory_config: Res<MemoryConfig>,
     memory_usage: Res<MemoryUsage>,
 ) where
-    T: Any + DataSize + Resource,
+    T: Any + Resource,
+    E: DataSizeEstimator<T>,
 {
     update_stats::<T, _>(memory_config, memory_usage, || {
-        MemoryStats::from_value(&*resource)
+        MemoryStats::from_value_with_estimator::<T, E>(&*resource)
     });
 }
 
-/// This system updates the [`MemoryStats`] for the given asset type `T`.
-pub fn update_stats_for_asset<T>(
+/// This system updates the [`MemoryStats`] for the given asset type `T`
+/// using the given [`DataSizeEstimator`].
+pub fn update_stats_for_asset<T, E>(
     assets: Res<Assets<T>>,
     memory_config: Res<MemoryConfig>,
     memory_usage: Res<MemoryUsage>,
 ) where
-    T: Any + DataSize + Asset,
+    T: Any + Asset,
+    E: DataSizeEstimator<T>,
 {
     update_stats::<T, _>(memory_config, memory_usage, || {
-        MemoryStats::from_values(assets.iter().map(|(_handle, asset)| asset))
+        MemoryStats::from_values_with_estimator::<T, E, _>(
+            assets.iter().map(|(_handle, asset)| asset),
+        )
     });
 }
 
 fn update_stats<T, F>(memory_config: Res<MemoryConfig>, memory_usage: Res<MemoryUsage>, op: F)
 where
-    T: Any + DataSize,
+    T: Any,
     F: FnOnce() -> MemoryStats,
 {
     if !memory_config.global {