Add MemoryStats and docs, docs, docs!

Author: BGR360

examples/components.rs

@@ -1,4 +1,5 @@
 use bevy::prelude::*;
+
 use bevy_datasize::{DataSize, MemoryUsage, MemoryUsagePlugin, RegisterSizedTypes};
 
 #[derive(Component, DataSize)]
@@ -24,7 +25,7 @@ fn spawn_entities(mut commands: Commands) {
 }
 
 fn print_size(memory_usage: Res<MemoryUsage>) {
-    let bytes = memory_usage.get_usage::<MyComponent>().unwrap();
+    let stats = memory_usage.get_stats::<MyComponent>().unwrap();
 
-    println!("Memory usage: {} bytes", bytes);
+    println!("Memory usage: {:#?}", stats);
 }

src/app_ext.rs

@@ -3,9 +3,15 @@ use std::any::Any;
 use bevy::app::App;
 use bevy::ecs::component::Component;
 
-use crate::{systems::update_datasize_for_component, DataSize, MemoryUsage};
+use crate::{systems::update_stats_for_component, DataSize, MemoryUsage};
 
+/// [`App`] extension methods to register types for memory usage tracking.
 pub trait RegisterSizedTypes {
+    /// Registers the given [`Component`] type with the
+    /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
+    ///
+    /// The given type `T` will be available to query on the [`MemoryUsage`]
+    /// resource.
     fn register_sized_component<T>(&mut self) -> &mut Self
     where
         T: Any + Component + DataSize;
@@ -22,7 +28,7 @@ impl RegisterSizedTypes for App {
 
         memory_usage.register_type::<T>();
 
-        self.add_system(update_datasize_for_component::<T>);
+        self.add_system(update_stats_for_component::<T>);
 
         self
     }

src/config.rs

@@ -1,12 +1,12 @@
+/// Configuration for the [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
 #[derive(Debug)]
 pub struct MemoryConfig {
-    pub tracking_enabled: bool,
+    /// Whether to track memory usage for all registered types.
+    pub global: bool,
 }
 
 impl Default for MemoryConfig {
     fn default() -> Self {
-        Self {
-            tracking_enabled: true,
-        }
+        Self { global: true }
     }
 }

src/lib.rs

@@ -1,24 +1,17 @@
 #![doc = include_str!("../README.md")]
+#![warn(missing_docs)]
 
 pub use datasize::DataSize;
 
 mod app_ext;
 mod config;
 mod plugin;
 mod resource;
-mod systems;
+mod stats;
+pub mod systems;
 
 pub use app_ext::RegisterSizedTypes;
 pub use config::MemoryConfig;
 pub use plugin::MemoryUsagePlugin;
 pub use resource::MemoryUsage;
-
-pub fn estimate_stack_and_heap_size<T>(value: &T) -> usize
-where
-    T: DataSize,
-{
-    let stack_size = std::mem::size_of::<T>();
-    let heap_size = value.estimate_heap_size();
-
-    stack_size + heap_size
-}
+pub use stats::MemoryStats;

src/plugin.rs

@@ -2,6 +2,9 @@ use bevy::app::Plugin;
 
 use crate::{MemoryConfig, MemoryUsage};
 
+/// Adds memory usage tracking to Apps.
+///
+/// You can configure this plugin using the [`MemoryConfig`] resource.
 pub struct MemoryUsagePlugin;
 
 impl Plugin for MemoryUsagePlugin {

src/resource.rs

@@ -2,22 +2,29 @@ use std::any::{Any, TypeId};
 
 use bevy::utils::HashMap;
 
+use crate::MemoryStats;
+
+/// Stores memory usage statistics for registered data types.
 #[derive(Debug, Default)]
 pub struct MemoryUsage {
-    pub(crate) datasizes: HashMap<TypeId, usize>,
+    pub(crate) datasizes: HashMap<TypeId, MemoryStats>,
 }
 
 impl MemoryUsage {
-    pub fn register_type<T>(&mut self)
+    /// Returns the most recent [`MemoryStats`] for the given type.
+    ///
+    /// Returns `None` if the type has not been registered.
+    pub fn get_stats<T>(&self) -> Option<MemoryStats>
     where
         T: Any,
     {
         let type_id = TypeId::of::<T>();
 
-        self.datasizes.insert(type_id, 0);
+        self.datasizes.get(&type_id).copied()
     }
 
-    pub fn update_usage<T>(&mut self, total_bytes: usize)
+    /// Updates the [`MemoryStats`] for the given type.
+    pub fn update_stats<T>(&mut self, stats: MemoryStats)
     where
         T: Any,
     {
@@ -28,15 +35,16 @@ impl MemoryUsage {
             .get_mut(&type_id)
             .expect("Memory usage not tracked for this type. Did you forget to register the type?");
 
-        *entry = total_bytes;
+        *entry = stats;
     }
 
-    pub fn get_usage<T>(&self) -> Option<usize>
+    /// Registers the given type with the usage tracker.
+    pub fn register_type<T>(&mut self)
     where
         T: Any,
     {
         let type_id = TypeId::of::<T>();
 
-        self.datasizes.get(&type_id).copied()
+        self.datasizes.insert(type_id, Default::default());
     }
 }

src/stats.rs

@@ -0,0 +1,109 @@
+use std::any::Any;
+
+use crate::DataSize;
+
+/// Memory usage statistics for a single data type.
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct MemoryStats {
+    /// The total number of instances of this data type.
+    pub count: usize,
+
+    /// The total number of "stack" bytes used by instances of this data type.
+    ///
+    /// See [`stack_size_of`] for details on the meaning of this quantity.
+    ///
+    /// [`stack_size_of`]: Self::stack_size_of
+    pub total_stack_bytes: usize,
+
+    /// The estimated total number of bytes used by instances of this data type.
+    ///
+    /// See [`heap_size_of`] for details on the meaning of this quantity.
+    ///
+    /// [`heap_size_of`]: Self::heap_size_of
+    pub total_heap_bytes: usize,
+}
+
+impl MemoryStats {
+    /// Returns the sum of `total_stack_bytes` and `total_heap_bytes` for
+    /// `self`.
+    #[inline]
+    pub fn total_bytes(&self) -> usize {
+        self.total_stack_bytes + self.total_heap_bytes
+    }
+
+    /// Returns the computed memory statistics for a single value.
+    #[inline]
+    pub fn from_value<T>(value: &T) -> Self
+    where
+        T: Any + DataSize,
+    {
+        Self {
+            count: 1,
+            total_stack_bytes: Self::stack_size_of(value),
+            total_heap_bytes: Self::heap_size_of(value),
+        }
+    }
+
+    /// Returns the computed memory statistics for a collection of values.
+    #[inline]
+    pub fn from_values<'a, T, I>(values: I) -> Self
+    where
+        T: Any + DataSize,
+        I: IntoIterator<Item = &'a T>,
+    {
+        let mut count = 0;
+        let mut total_stack_bytes = 0;
+        let mut total_heap_bytes = 0;
+
+        for value in values.into_iter() {
+            count += 1;
+            total_stack_bytes += Self::stack_size_of(value);
+            total_heap_bytes += Self::heap_size_of(value);
+        }
+
+        Self {
+            count,
+            total_stack_bytes,
+            total_heap_bytes,
+        }
+    }
+
+    /// Returns the "stack" size of the given value.
+    ///
+    /// This quantity represents how many bytes the type *would* take up on the
+    /// stack if were allocated there.
+    ///
+    /// This quantity is **exact** and is computed using [`std::mem::size_of`].
+    #[inline]
+    pub fn stack_size_of<T>(value: &T) -> usize
+    where
+        T: Any,
+    {
+        std::mem::size_of_val(value)
+    }
+
+    /// Returns the estimated heap size of the given value.
+    ///
+    /// This quantity represents how many bytes the type occupies apart from the
+    /// immediate bytes of its fields.
+    #[inline]
+    pub fn heap_size_of<T>(value: &T) -> usize
+    where
+        T: DataSize,
+    {
+        value.estimate_heap_size()
+    }
+
+    /// 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
+    where
+        T: Any + DataSize,
+    {
+        Self::stack_size_of(value) + Self::heap_size_of(value)
+    }
+}

src/systems.rs

@@ -1,22 +1,26 @@
+//! Systems used by this library.
+
 use std::any::Any;
 
 use bevy::prelude::*;
 use datasize::DataSize;
 
-use crate::{estimate_stack_and_heap_size, MemoryConfig, MemoryUsage};
+use crate::{MemoryConfig, MemoryStats, MemoryUsage};
 
-pub fn update_datasize_for_component<T>(
+/// This system iterates through all components of type `T` and updates their
+/// [`MemoryStats`].
+pub fn update_stats_for_component<T>(
     query: Query<&T>,
     memory_config: Res<MemoryConfig>,
     mut memory_usage: ResMut<MemoryUsage>,
 ) where
     T: Any + Component + DataSize,
 {
-    if !memory_config.tracking_enabled {
+    if !memory_config.global {
         return;
     }
 
-    let total_bytes: usize = query.iter().map(estimate_stack_and_heap_size).sum();
+    let stats = MemoryStats::from_values(query.iter());
 
-    memory_usage.update_usage::<T>(total_bytes);
+    memory_usage.update_stats::<T>(stats);
 }