From 1d4176d4cd5ae658f7cf2eea3da427622a59b319 Mon Sep 17 00:00:00 2001 From: Mateusz Wachowiak Date: Wed, 3 Apr 2024 04:47:08 +0200 Subject: [PATCH] Add methods `iter_resources` and `iter_resources_mut` (#12829) # Objective - Closes #12019 - Related to #4955 - Useful for dev_tools and networking ## Solution - Create `World::iter_resources()` and `World::iter_resources_mut()` --------- Co-authored-by: Alice Cecile Co-authored-by: James Liu Co-authored-by: Pablo Reinhardt <126117294+pablo-lua@users.noreply.github.com> --- crates/bevy_ecs/src/world/mod.rs | 269 +++++++++++++++++++++++++++++++ 1 file changed, 269 insertions(+) diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index d245df8241930..ed6781ffe3596 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -2136,6 +2136,209 @@ impl World { } } + /// Iterates over all resources in the world. + /// + /// The returned iterator provides lifetimed, but type-unsafe pointers. Actually reading the contents + /// of each resource will require the use of unsafe code. + /// + /// # Examples + /// + /// ## Printing the size of all resources + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # #[derive(Resource)] + /// # struct A(u32); + /// # #[derive(Resource)] + /// # struct B(u32); + /// # + /// # let mut world = World::new(); + /// # world.insert_resource(A(1)); + /// # world.insert_resource(B(2)); + /// let mut total = 0; + /// for (info, _) in world.iter_resources() { + /// println!("Resource: {}", info.name()); + /// println!("Size: {} bytes", info.layout().size()); + /// total += info.layout().size(); + /// } + /// println!("Total size: {} bytes", total); + /// # assert_eq!(total, std::mem::size_of::() + std::mem::size_of::()); + /// ``` + /// + /// ## Dynamically running closures for resources matching specific `TypeId`s + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # use std::collections::HashMap; + /// # use std::any::TypeId; + /// # use bevy_ptr::Ptr; + /// # #[derive(Resource)] + /// # struct A(u32); + /// # #[derive(Resource)] + /// # struct B(u32); + /// # + /// # let mut world = World::new(); + /// # world.insert_resource(A(1)); + /// # world.insert_resource(B(2)); + /// # + /// // In this example, `A` and `B` are resources. We deliberately do not use the + /// // `bevy_reflect` crate here to showcase the low-level [`Ptr`] usage. You should + /// // probably use something like `ReflectFromPtr` in a real-world scenario. + /// + /// // Create the hash map that will store the closures for each resource type + /// let mut closures: HashMap)>> = HashMap::new(); + /// + /// // Add closure for `A` + /// closures.insert(TypeId::of::(), Box::new(|ptr| { + /// // SAFETY: We assert ptr is the same type of A with TypeId of A + /// let a = unsafe { &ptr.deref::() }; + /// # assert_eq!(a.0, 1); + /// // ... do something with `a` here + /// })); + /// + /// // Add closure for `B` + /// closures.insert(TypeId::of::(), Box::new(|ptr| { + /// // SAFETY: We assert ptr is the same type of B with TypeId of B + /// let b = unsafe { &ptr.deref::() }; + /// # assert_eq!(b.0, 2); + /// // ... do something with `b` here + /// })); + /// + /// // Iterate all resources, in order to run the closures for each matching resource type + /// for (info, ptr) in world.iter_resources() { + /// let Some(type_id) = info.type_id() else { + /// // It's possible for resources to not have a `TypeId` (e.g. non-Rust resources + /// // dynamically inserted via a scripting language) in which case we can't match them. + /// continue; + /// }; + /// + /// let Some(closure) = closures.get(&type_id) else { + /// // No closure for this resource type, skip it. + /// continue; + /// }; + /// + /// // Run the closure for the resource + /// closure(&ptr); + /// } + /// ``` + #[inline] + pub fn iter_resources(&self) -> impl Iterator)> { + self.storages + .resources + .iter() + .filter_map(|(component_id, data)| { + // SAFETY: If a resource has been initialized, a corresponding ComponentInfo must exist with it's ID. + let component_info = unsafe { + self.components + .get_info(component_id) + .debug_checked_unwrap() + }; + Some((component_info, data.get_data()?)) + }) + } + + /// Mutably iterates over all resources in the world. + /// + /// The returned iterator provides lifetimed, but type-unsafe pointers. Actually reading from or writing + /// to the contents of each resource will require the use of unsafe code. + /// + /// # Example + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # use bevy_ecs::change_detection::MutUntyped; + /// # use std::collections::HashMap; + /// # use std::any::TypeId; + /// # #[derive(Resource)] + /// # struct A(u32); + /// # #[derive(Resource)] + /// # struct B(u32); + /// # + /// # let mut world = World::new(); + /// # world.insert_resource(A(1)); + /// # world.insert_resource(B(2)); + /// # + /// // In this example, `A` and `B` are resources. We deliberately do not use the + /// // `bevy_reflect` crate here to showcase the low-level `MutUntyped` usage. You should + /// // probably use something like `ReflectFromPtr` in a real-world scenario. + /// + /// // Create the hash map that will store the mutator closures for each resource type + /// let mut mutators: HashMap)>> = HashMap::new(); + /// + /// // Add mutator closure for `A` + /// mutators.insert(TypeId::of::(), Box::new(|mut_untyped| { + /// // Note: `MutUntyped::as_mut()` automatically marks the resource as changed + /// // for ECS change detection, and gives us a `PtrMut` we can use to mutate the resource. + /// // SAFETY: We assert ptr is the same type of A with TypeId of A + /// let a = unsafe { &mut mut_untyped.as_mut().deref_mut::() }; + /// # a.0 += 1; + /// // ... mutate `a` here + /// })); + /// + /// // Add mutator closure for `B` + /// mutators.insert(TypeId::of::(), Box::new(|mut_untyped| { + /// // SAFETY: We assert ptr is the same type of B with TypeId of B + /// let b = unsafe { &mut mut_untyped.as_mut().deref_mut::() }; + /// # b.0 += 1; + /// // ... mutate `b` here + /// })); + /// + /// // Iterate all resources, in order to run the mutator closures for each matching resource type + /// for (info, mut mut_untyped) in world.iter_resources_mut() { + /// let Some(type_id) = info.type_id() else { + /// // It's possible for resources to not have a `TypeId` (e.g. non-Rust resources + /// // dynamically inserted via a scripting language) in which case we can't match them. + /// continue; + /// }; + /// + /// let Some(mutator) = mutators.get(&type_id) else { + /// // No mutator closure for this resource type, skip it. + /// continue; + /// }; + /// + /// // Run the mutator closure for the resource + /// mutator(&mut mut_untyped); + /// } + /// # assert_eq!(world.resource::().0, 2); + /// # assert_eq!(world.resource::().0, 3); + /// ``` + #[inline] + pub fn iter_resources_mut(&mut self) -> impl Iterator)> { + self.storages + .resources + .iter() + .filter_map(|(component_id, data)| { + // SAFETY: If a resource has been initialized, a corresponding ComponentInfo must exist with it's ID. + let component_info = unsafe { + self.components + .get_info(component_id) + .debug_checked_unwrap() + }; + let (ptr, ticks) = data.get_with_ticks()?; + + // SAFETY: + // - We have exclusive access to the world, so no other code can be aliasing the `TickCells` + // - We only hold one `TicksMut` at a time, and we let go of it before getting the next one + let ticks = unsafe { + TicksMut::from_tick_cells( + ticks, + self.last_change_tick(), + self.read_change_tick(), + ) + }; + + let mut_untyped = MutUntyped { + // SAFETY: + // - We have exclusive access to the world, so no other code can be aliasing the `Ptr` + // - We iterate one resource at a time, and we let go of each `PtrMut` before getting the next one + value: unsafe { ptr.assert_unique() }, + ticks, + }; + + Some((component_info, mut_untyped)) + }) + } + /// Gets a `!Send` resource to the resource with the id [`ComponentId`] if it exists. /// The returned pointer must not be used to modify the resource, and must not be /// dereferenced after the immutable borrow of the [`World`] ends. @@ -2554,6 +2757,12 @@ mod tests { #[derive(Resource)] struct TestResource(u32); + #[derive(Resource)] + struct TestResource2(String); + + #[derive(Resource)] + struct TestResource3; + #[test] fn get_resource_by_id() { let mut world = World::new(); @@ -2594,6 +2803,66 @@ mod tests { assert_eq!(resource.0, 43); } + #[test] + fn iter_resources() { + let mut world = World::new(); + world.insert_resource(TestResource(42)); + world.insert_resource(TestResource2("Hello, world!".to_string())); + world.insert_resource(TestResource3); + world.remove_resource::(); + + let mut iter = world.iter_resources(); + + let (info, ptr) = iter.next().unwrap(); + assert_eq!(info.name(), std::any::type_name::()); + // SAFETY: We know that the resource is of type `TestResource` + assert_eq!(unsafe { ptr.deref::().0 }, 42); + + let (info, ptr) = iter.next().unwrap(); + assert_eq!(info.name(), std::any::type_name::()); + assert_eq!( + // SAFETY: We know that the resource is of type `TestResource2` + unsafe { &ptr.deref::().0 }, + &"Hello, world!".to_string() + ); + + assert!(iter.next().is_none()); + } + + #[test] + fn iter_resources_mut() { + let mut world = World::new(); + world.insert_resource(TestResource(42)); + world.insert_resource(TestResource2("Hello, world!".to_string())); + world.insert_resource(TestResource3); + world.remove_resource::(); + + let mut iter = world.iter_resources_mut(); + + let (info, mut mut_untyped) = iter.next().unwrap(); + assert_eq!(info.name(), std::any::type_name::()); + // SAFETY: We know that the resource is of type `TestResource` + unsafe { + mut_untyped.as_mut().deref_mut::().0 = 43; + }; + + let (info, mut mut_untyped) = iter.next().unwrap(); + assert_eq!(info.name(), std::any::type_name::()); + // SAFETY: We know that the resource is of type `TestResource2` + unsafe { + mut_untyped.as_mut().deref_mut::().0 = "Hello, world?".to_string(); + }; + + assert!(iter.next().is_none()); + std::mem::drop(iter); + + assert_eq!(world.resource::().0, 43); + assert_eq!( + world.resource::().0, + "Hello, world?".to_string() + ); + } + #[test] fn custom_resource_with_layout() { static DROP_COUNT: AtomicU32 = AtomicU32::new(0);