diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index e4c45d2b69081..7c00fd6e6f1db 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -2184,6 +2184,194 @@ impl World { .get_mut_by_id(component_id) } } + + /// Iterates over all resources in the world + /// + /// # 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!("----------------"); + /// 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| { + /// 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| { + /// 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); + /// } + /// ``` + pub fn iter_resources(&self) -> impl Iterator)> { + self.storages.resources.iter().map(|(component_id, data)| { + let component_info = self.components.get_info(component_id).unwrap_or_else(|| { + panic!("ComponentInfo should exist for all resources in the world, but it does not for {:?}", component_id); + }); + let ptr = data.get_data().unwrap_or_else(|| { + panic!( + "When iterating all resources, resource of type {} was supposed to exist, but did not.", + component_info.name() + ) + }); + (component_info, ptr) + }) + } + + /// Mutably iterates over all resources in the world + /// + /// # 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. + /// 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| { + /// 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); + /// ``` + pub fn iter_resources_mut(&mut self) -> impl Iterator)> { + self.storages.resources.iter().map(|(component_id, data)| { + let component_info = self.components.get_info(component_id).unwrap_or_else(|| { + panic!("ComponentInfo should exist for all resources in the world, but it does not for {:?}", component_id); + }); + + let (ptr, ticks) = data.get_with_ticks().unwrap_or_else(|| { + panic!( + "When iterating all resources, resource of type {} was supposed to exist, but did not.", + component_info.name() + ) + }); + + // 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, + }; + + (component_info, mut_untyped) + }) + } } // Schedule-related methods @@ -2495,6 +2683,9 @@ mod tests { #[derive(Resource)] struct TestResource(u32); + #[derive(Resource)] + struct TestResource2(String); + #[test] fn get_resource_by_id() { let mut world = World::new(); @@ -2796,4 +2987,63 @@ mod tests { let mut world = World::new(); world.spawn(()); } + + #[test] + fn iterate_resources() { + let mut world = World::new(); + + world.insert_resource(TestResource(42)); + world.insert_resource(TestResource2("Hello, world!".to_string())); + + 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 iterate_resources_mut() { + let mut world = World::new(); + + world.insert_resource(TestResource(42)); + world.insert_resource(TestResource2("Hello, world!".to_string())); + + 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() + ); + } }