Initial working version.

Author: BGR360

README.md

@@ -0,0 +1,5 @@
+# bevy_datasize
+
+This is a library for tracking memory usage in [Bevy] apps.
+
+[Bevy]: https://lib.rs/bevy

examples/components.rs

@@ -1,13 +1,30 @@
 use bevy::prelude::*;
-use bevy_datasize::DataSize;
+use bevy_datasize::{DataSize, MemoryUsage, MemoryUsagePlugin, RegisterSizedTypes};
 
 #[derive(Component, DataSize)]
 struct MyComponent {
     pub vec: Vec<u8>,
 }
 
 fn main() {
-    let component = MyComponent { vec: vec![0; 4096] };
+    App::new()
+        .add_plugins(MinimalPlugins)
+        .add_plugin(MemoryUsagePlugin)
+        .register_sized_component::<MyComponent>()
+        .add_startup_system(spawn_entities)
+        .add_system(print_size)
+        .run();
+}
+
+fn spawn_entities(mut commands: Commands) {
+    for _ in 0..8 {
+        let component = MyComponent { vec: vec![0; 1024] };
+        commands.spawn().insert(component);
+    }
+}
+
+fn print_size(memory_usage: Res<MemoryUsage>) {
+    let bytes = memory_usage.get_usage::<MyComponent>().unwrap();
 
-    println!("{}", component.estimate_heap_size());
+    println!("Memory usage: {} bytes", bytes);
 }

src/app_ext.rs

@@ -0,0 +1,29 @@
+use std::any::Any;
+
+use bevy::app::App;
+use bevy::ecs::component::Component;
+
+use crate::{systems::update_datasize_for_component, DataSize, MemoryUsage};
+
+pub trait RegisterSizedTypes {
+    fn register_sized_component<T>(&mut self) -> &mut Self
+    where
+        T: Any + Component + DataSize;
+}
+
+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`?",
+        );
+
+        memory_usage.register_type::<T>();
+
+        self.add_system(update_datasize_for_component::<T>);
+
+        self
+    }
+}

src/config.rs

@@ -0,0 +1,12 @@
+#[derive(Debug)]
+pub struct MemoryConfig {
+    pub tracking_enabled: bool,
+}
+
+impl Default for MemoryConfig {
+    fn default() -> Self {
+        Self {
+            tracking_enabled: true,
+        }
+    }
+}

src/lib.rs

@@ -1,4 +1,24 @@
 #![doc = include_str!("../README.md")]
-#![warn(missing_docs)]
 
 pub use datasize::DataSize;
+
+mod app_ext;
+mod config;
+mod plugin;
+mod resource;
+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
+}

src/plugin.rs

@@ -0,0 +1,12 @@
+use bevy::app::Plugin;
+
+use crate::{MemoryConfig, MemoryUsage};
+
+pub struct MemoryUsagePlugin;
+
+impl Plugin for MemoryUsagePlugin {
+    fn build(&self, app: &mut bevy::prelude::App) {
+        app.init_resource::<MemoryConfig>();
+        app.init_resource::<MemoryUsage>();
+    }
+}

src/resource.rs

@@ -0,0 +1,42 @@
+use std::any::{Any, TypeId};
+
+use bevy::utils::HashMap;
+
+#[derive(Debug, Default)]
+pub struct MemoryUsage {
+    pub(crate) datasizes: HashMap<TypeId, usize>,
+}
+
+impl MemoryUsage {
+    pub fn register_type<T>(&mut self)
+    where
+        T: Any,
+    {
+        let type_id = TypeId::of::<T>();
+
+        self.datasizes.insert(type_id, 0);
+    }
+
+    pub fn update_usage<T>(&mut self, total_bytes: usize)
+    where
+        T: Any,
+    {
+        let type_id = TypeId::of::<T>();
+
+        let entry = self
+            .datasizes
+            .get_mut(&type_id)
+            .expect("Memory usage not tracked for this type. Did you forget to register the type?");
+
+        *entry = total_bytes;
+    }
+
+    pub fn get_usage<T>(&self) -> Option<usize>
+    where
+        T: Any,
+    {
+        let type_id = TypeId::of::<T>();
+
+        self.datasizes.get(&type_id).copied()
+    }
+}

src/systems.rs

@@ -0,0 +1,22 @@
+use std::any::Any;
+
+use bevy::prelude::*;
+use datasize::DataSize;
+
+use crate::{estimate_stack_and_heap_size, MemoryConfig, MemoryUsage};
+
+pub fn update_datasize_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 {
+        return;
+    }
+
+    let total_bytes: usize = query.iter().map(estimate_stack_and_heap_size).sum();
+
+    memory_usage.update_usage::<T>(total_bytes);
+}