Skip to content

Add methods iter_resources and iter_resources_mut #12829

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Apr 3, 2024
Merged

Add methods iter_resources and iter_resources_mut #12829

Changes from 3 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b1c4813
add iter_resources and iter_resources_mut methods on world
matiqo15 Apr 1, 2024
058884b
add tests for iter_resources and iter_resources_mut
matiqo15 Apr 1, 2024
cb555a0
move safety comment
matiqo15 Apr 1, 2024
0174d38
add doc link
matiqo15 Apr 1, 2024
4f2feda
improve docs
matiqo15 Apr 2, 2024
009535d
use `debug_checked_unwrap`
matiqo15 Apr 2, 2024
99a0b46
format
matiqo15 Apr 2, 2024
03d64d8
safety comments
matiqo15 Apr 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
250 changes: 250 additions & 0 deletions crates/bevy_ecs/src/world/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2136,6 +2136,190 @@ impl World {
}
}

/// 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::<A>() + std::mem::size_of::<B>());
/// ```
///
/// ## 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<TypeId, Box<dyn Fn(&Ptr<'_>)>> = HashMap::new();
///
/// // Add closure for `A`
/// closures.insert(TypeId::of::<A>(), Box::new(|ptr| {
/// let a = unsafe { &ptr.deref::<A>() };
/// # assert_eq!(a.0, 1);
/// // ... do something with `a` here
/// }));
///
/// // Add closure for `B`
/// closures.insert(TypeId::of::<B>(), Box::new(|ptr| {
/// let b = unsafe { &ptr.deref::<B>() };
/// # 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<Item = (&ComponentInfo, Ptr<'_>)> {
self.storages
.resources
.iter()
.filter_map(|(component_id, data)| {
let component_info = self.components.get_info(component_id)?;
Some((component_info, data.get_data()?))
})
}

/// 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<TypeId, Box<dyn Fn(&mut MutUntyped<'_>)>> = HashMap::new();
///
/// // Add mutator closure for `A`
/// mutators.insert(TypeId::of::<A>(), 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>() };
/// # a.0 += 1;
/// // ... mutate `a` here
/// }));
///
/// // Add mutator closure for `B`
/// mutators.insert(TypeId::of::<B>(), Box::new(|mut_untyped| {
/// let b = unsafe { &mut mut_untyped.as_mut().deref_mut::<B>() };
/// # 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::<A>().0, 2);
/// # assert_eq!(world.resource::<B>().0, 3);
/// ```
#[inline]
pub fn iter_resources_mut(&mut self) -> impl Iterator<Item = (&ComponentInfo, MutUntyped<'_>)> {
self.storages
.resources
.iter()
.filter_map(|(component_id, data)| {
let component_info = self.components.get_info(component_id)?;
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 oned
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.
Expand Down Expand Up @@ -2554,6 +2738,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();
Expand Down Expand Up @@ -2594,6 +2784,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::<TestResource3>();

let mut iter = world.iter_resources();

let (info, ptr) = iter.next().unwrap();
assert_eq!(info.name(), std::any::type_name::<TestResource>());
// SAFETY: We know that the resource is of type `TestResource`
assert_eq!(unsafe { ptr.deref::<TestResource>().0 }, 42);

let (info, ptr) = iter.next().unwrap();
assert_eq!(info.name(), std::any::type_name::<TestResource2>());
assert_eq!(
// SAFETY: We know that the resource is of type `TestResource2`
unsafe { &ptr.deref::<TestResource2>().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::<TestResource3>();

let mut iter = world.iter_resources_mut();

let (info, mut mut_untyped) = iter.next().unwrap();
assert_eq!(info.name(), std::any::type_name::<TestResource>());
// SAFETY: We know that the resource is of type `TestResource`
unsafe {
mut_untyped.as_mut().deref_mut::<TestResource>().0 = 43;
};

let (info, mut mut_untyped) = iter.next().unwrap();
assert_eq!(info.name(), std::any::type_name::<TestResource2>());
// SAFETY: We know that the resource is of type `TestResource2`
unsafe {
mut_untyped.as_mut().deref_mut::<TestResource2>().0 = "Hello, world?".to_string();
};

assert!(iter.next().is_none());
std::mem::drop(iter);

assert_eq!(world.resource::<TestResource>().0, 43);
assert_eq!(
world.resource::<TestResource2>().0,
"Hello, world?".to_string()
);
}

#[test]
fn custom_resource_with_layout() {
static DROP_COUNT: AtomicU32 = AtomicU32::new(0);
Expand Down