Add support for resources and assets, and boof up docs.

Author: BGR360

examples/basic.rs

@@ -0,0 +1,63 @@
+//! This example shows how to register your own resource, component, and asset
+//! types for memory usage tracking.
+
+use bevy::{prelude::*, reflect::TypeUuid};
+use bevy_datasize::prelude::*;
+
+#[derive(DataSize)]
+struct MyResource {
+    pub data: Vec<u8>,
+}
+
+#[derive(Component, DataSize)]
+struct MyComponent {
+    pub data: Vec<u8>,
+}
+
+#[derive(Debug, TypeUuid)]
+#[uuid = "0669d1a3-34b5-4bac-b9cd-26a2e2df79ff"]
+pub struct MyAsset {
+    pub data: Vec<u8>,
+}
+
+fn main() {
+    App::new()
+        .add_plugins(MinimalPlugins)
+        .add_plugin(MemoryUsagePlugin)
+        .register_sized_resource::<MyResource>()
+        .register_sized_component::<MyComponent>()
+        .insert_resource(MyResource {
+            data: vec![42; 4096],
+        })
+        .add_startup_system(spawn_entities)
+        .add_startup_system(add_asset)
+        .add_system(print_sizes)
+        .run();
+}
+
+fn spawn_entities(mut commands: Commands) {
+    for _ in 0..8 {
+        let component = MyComponent {
+            data: vec![0; 1024],
+        };
+        commands.spawn().insert(component);
+    }
+}
+
+fn add_asset(mut handle: Local<Handle<MyAsset>>, mut assets: ResMut<Assets<MyAsset>>) {
+    *handle = assets.add(MyAsset {
+        data: vec![68; 425],
+    });
+}
+
+fn print_sizes(memory_usage: Res<MemoryUsage>) {
+    let resource_stats = memory_usage.get_stats::<MyResource>().unwrap();
+    let component_stats = memory_usage.get_stats::<MyComponent>().unwrap();
+    let asset_stats = memory_usage.get_stats::<MyAsset>().unwrap();
+
+    println!();
+    println!("Memory usage:");
+    println!("MyResource: {:#?}", resource_stats);
+    println!("MyComponent: {:#?}", component_stats);
+    println!("MyAsset: {:#?}", asset_stats);
+}

examples/components.rs

@@ -1,9 +1,12 @@
 use std::any::Any;
 
-use bevy::app::App;
-use bevy::ecs::component::Component;
+use bevy::{
+    app::App,
+    asset::Asset,
+    ecs::{component::Component, system::Resource},
+};
 
-use crate::{systems::update_stats_for_component, DataSize, MemoryUsage};
+use crate::{systems, DataSize, MemoryUsage};
 
 /// [`App`] extension methods to register types for memory usage tracking.
 pub trait RegisterSizedTypes {
@@ -14,22 +17,69 @@ pub trait RegisterSizedTypes {
     /// resource.
     fn register_sized_component<T>(&mut self) -> &mut Self
     where
-        T: Any + Component + DataSize;
+        T: Any + DataSize + Component;
+
+    /// Registers the given [`Resource`] type with the
+    /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
+    ///
+    /// The given type `T` will be available to query on the [`MemoryUsage`]
+    /// resource.
+    fn register_sized_resource<T>(&mut self) -> &mut Self
+    where
+        T: Any + DataSize + Resource;
+
+    /// Registers the given [`Asset`] type with the
+    /// [`MemoryUsagePlugin`][crate::MemoryUsagePlugin].
+    ///
+    /// The given type `T` will be available to query on the [`MemoryUsage`]
+    /// 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,
     {
-        let mut memory_usage = self.world.get_resource_mut::<MemoryUsage>().expect(
-            "Cannot find resource `MemoryUsage`. Did you forget to add the `MemoryUsagePlugin`?",
-        );
+        register_type::<T>(self);
+
+        self.add_system(systems::update_stats_for_component::<T>);
+
+        self
+    }
+
+    fn register_sized_resource<T>(&mut self) -> &mut Self
+    where
+        T: Any + DataSize + Resource,
+    {
+        register_type::<T>(self);
+
+        self.add_system(systems::update_stats_for_resource::<T>);
+
+        self
+    }
 
-        memory_usage.register_type::<T>();
+    fn register_sized_asset<T>(&mut self) -> &mut Self
+    where
+        T: Any + DataSize + Asset,
+    {
+        register_type::<T>(self);
 
-        self.add_system(update_stats_for_component::<T>);
+        self.add_system(systems::update_stats_for_asset::<T>);
 
         self
     }
 }
+
+fn register_type<T>(app: &mut App)
+where
+    T: Any + DataSize,
+{
+    let mut memory_usage = app.world.get_resource_mut::<MemoryUsage>().expect(
+        "Cannot find resource `MemoryUsage`. Did you forget to add the `MemoryUsagePlugin`?",
+    );
+
+    memory_usage.register_type::<T>();
+}

src/app_ext.rs

@@ -1,4 +1,42 @@
-#![doc = include_str!("../README.md")]
+//! This is a library for tracking memory usage in [Bevy](https://lib.rs/bevy) apps.
+//!
+//! It is based on the [`datasize`] crate.
+//!
+//! # Example
+//!
+//! ```no_run
+//! use bevy::prelude::*;
+//! use bevy_datasize::prelude::*;
+//!
+//! #[derive(Component, DataSize)]
+//! struct MyComponent {
+//!     data: Vec<u8>,
+//! }
+//!
+//! fn main() {
+//!     App::new()
+//!         .add_plugins(MinimalPlugins)
+//!         .add_plugin(MemoryUsagePlugin)
+//!         .register_sized_component::<MyComponent>()
+//!         .add_system(print_datasize)
+//!         .run();
+//! }
+//!
+//! fn print_datasize(memory_usage: Res<MemoryUsage>) {
+//!     let MemoryStats {
+//!         count,
+//!         total_stack_bytes,
+//!         total_heap_bytes,
+//!     } = memory_usage.get_stats::<MyComponent>().unwrap();
+//!
+//!     println!("MyComponent count: {count}");
+//!     println!("MyComponent total stack usage: {total_stack_bytes} bytes");
+//!     println!("MyComponent total heap usage: {total_heap_bytes} bytes");
+//! }
+//! ```
+//!
+//! See the [`datasize`] docs for more information on the [`DataSize`] trait.
+
 #![warn(missing_docs)]
 
 pub use datasize::DataSize;
@@ -18,5 +56,12 @@ pub use stats::MemoryStats;
 
 #[allow(missing_docs)]
 pub mod prelude {
-    pub use crate::{MemoryConfig, MemoryUsage, MemoryUsagePlugin, RegisterSizedTypes};
+    pub use crate::{
+        DataSize, MemoryConfig, MemoryStats, MemoryUsage, MemoryUsagePlugin, RegisterSizedTypes,
+    };
 }
+
+// Make sure the README example compiles
+#[cfg(doctest)]
+#[doc = include_str!("../README.md")]
+pub struct ReadmeTest;

src/lib.rs

@@ -2,7 +2,13 @@
 
 use std::any::Any;
 
-use bevy::prelude::*;
+use bevy::{
+    asset::{Asset, Assets},
+    ecs::{
+        component::Component,
+        system::{Query, Res, Resource},
+    },
+};
 use datasize::DataSize;
 
 use crate::{MemoryConfig, MemoryStats, MemoryUsage};
@@ -15,12 +21,48 @@ pub fn update_stats_for_component<T>(
     memory_usage: Res<MemoryUsage>,
 ) where
     T: Any + Component + DataSize,
+{
+    update_stats::<T, _>(memory_config, memory_usage, || {
+        MemoryStats::from_values(query.iter())
+    });
+}
+
+/// This system updates the [`MemoryStats`] for the given resource type `T`.
+pub fn update_stats_for_resource<T>(
+    resource: Res<T>,
+    memory_config: Res<MemoryConfig>,
+    memory_usage: Res<MemoryUsage>,
+) where
+    T: Any + DataSize + Resource,
+{
+    update_stats::<T, _>(memory_config, memory_usage, || {
+        MemoryStats::from_value(&*resource)
+    });
+}
+
+/// This system updates the [`MemoryStats`] for the given asset type `T`.
+pub fn update_stats_for_asset<T>(
+    assets: Res<Assets<T>>,
+    memory_config: Res<MemoryConfig>,
+    memory_usage: Res<MemoryUsage>,
+) where
+    T: Any + DataSize + Asset,
+{
+    update_stats::<T, _>(memory_config, memory_usage, || {
+        MemoryStats::from_values(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,
+    F: FnOnce() -> MemoryStats,
 {
     if !memory_config.global {
         return;
     }
 
-    let stats = MemoryStats::from_values(query.iter());
+    let stats = op();
 
     memory_usage.update_stats_fast::<T>(stats);
 }