Fetch multiple resources at once from World

crates/bevy_ecs/src/lib.rs

@@ -16,6 +16,7 @@ pub mod query;
 #[cfg(feature = "bevy_reflect")]
 pub mod reflect;
 pub mod removal_detection;
+pub mod resource_bundle;
 pub mod schedule;
 pub mod storage;
 pub mod system;
@@ -1403,6 +1404,49 @@ mod tests {
         query.iter(&world_b).for_each(|_| {});
     }
 
+    #[derive(Resource)]
+    struct Num(isize);
+
+    #[derive(Resource)]
+    struct BigNum(i128);
+
+    #[derive(Resource)]
+    struct SmallNum(i8);
+
+    #[test]
+    fn resource_bundle() {
+        let mut world = World::new();
+        world.insert_resource(SmallNum(1));
+        world.insert_resource(Num(1_000));
+        world.insert_resource(BigNum(1_000_000));
+
+        let (small_num, mut normal_num, big_num) = world
+            .get_resources_mut::<(SmallNum, Num, BigNum)>()
+            .expect("Couldn't get one of the resources in the bundle.");
+
+        assert_eq!(small_num.0, 1);
+        assert_eq!(big_num.0, 1_000_000);
+
+        normal_num.0 *= 2;
+
+        let normal_num = world.get_resource::<Num>().unwrap();
+
+        assert_eq!(normal_num.0, 2_000);
+    }
+
+    #[test]
+    #[should_panic]
+    #[cfg(not(miri))]
+    fn access_conflict_in_resource_bundle() {
+        let mut world = World::new();
+        world.insert_resource(SmallNum(1));
+        world.insert_resource(Num(1_000));
+
+        let (_small, _num, _num_mut) = world
+            .get_resources_mut::<(SmallNum, Num, Num)>() // This is an access conflict!
+            .unwrap();
+    }
+
     #[test]
     fn resource_scope() {
         let mut world = World::default();

crates/bevy_ecs/src/resource_bundle.rs

@@ -0,0 +1,148 @@
+//! This module contains the logic for bundling up resources together.
+use bevy_utils::{all_tuples, TypeIdSet};
+use std::any::TypeId;
+
+use crate::{
+    prelude::Mut,
+    system::{Res, Resource},
+    world::unsafe_world_cell::UnsafeWorldCell,
+};
+
+/// Bundle of resources. With this trait we can fetch multiple resources at once from a world.
+pub trait ResourceBundle {
+    /// Write access to the resources of this resource bundle. This type should provide write access, like `&mut R` or `ResMut<R>`
+    type WriteAccess<'a>;
+    /// Read-only access to the resources of this resource bundle. This type should provide read-only access, like `&R` or `Res<R>`
+    type ReadOnlyAccess<'a>;
+    /// Get write access to the resources in the bundle.
+    ///
+    /// # Safety
+    /// The caller must ensure that each resource in this bundle is safe to access mutably.
+    /// For example, if `R` is in the bundle, there should not be any other valid references to R.
+    unsafe fn fetch_write_access(world: UnsafeWorldCell<'_>) -> Option<Self::WriteAccess<'_>>;
+    /// Get read-only access to the resources in this bundle.
+    ///
+    /// # Safety
+    /// The caller must it is valid to get read-only access to each of the resources in this bundle.
+    /// For example, if `R` is in the bundle, there should not be any valid *mutable* references to R.
+    unsafe fn fetch_read_only(world: UnsafeWorldCell<'_>) -> Option<Self::ReadOnlyAccess<'_>>;
+    /// Return `true` if there are access conflicts within the bundle. In other words, this returns `true`
+    /// if and only a resource appears twice in the bundle.
+    fn contains_access_conflicts() -> bool {
+        false
+    }
+}
+
+/// This isn't public and part of the [`ResourceBundle`] trait because [`BundleAccessTable`] shouldn't be public.
+trait AccessConflictTracker {
+    /// Merge the internal [`access table`](BundleAccessTable) with some external one.
+    fn merge_with(other: &mut BundleAccessTable);
+    /// Return `true` if there is conflicting access within the bundle. For example, two mutable references
+    /// to the same resource.
+    fn contains_conflicting_access() -> bool {
+        false
+    }
+}
+
+struct BundleAccessTable {
+    table: TypeIdSet,
+    conflicted: bool,
+}
+
+impl BundleAccessTable {
+    /// Create a new empty access table.
+    fn new_empty_unconflicted() -> Self {
+        Self {
+            table: TypeIdSet::default(),
+            conflicted: false,
+        }
+    }
+
+    /// Insert a key-value pair to the table. If the insert causes an access conflict,
+    /// the internal conflict flag will be set to `true`.
+    /// # NOTE
+    /// Even if the insertion solved an existing conflict, this will not be reflected in the internal conflict flag.
+    fn insert_checked(&mut self, id: TypeId) {
+        self.conflicted |= !self.table.insert(id);
+    }
+
+    /// Returns the internal access conflict flag.
+    /// If this is `true`, that means that either the internal table contains an access conflict,
+    /// or at one point there was an attempt to merge this table with a conflicted one.
+    fn is_conflicted(&self) -> bool {
+        self.conflicted
+    }
+}
+
+impl<R: Resource> ResourceBundle for R {
+    type WriteAccess<'a> = Mut<'a, R>;
+    type ReadOnlyAccess<'a> = &'a R;
+    unsafe fn fetch_write_access(world: UnsafeWorldCell<'_>) -> Option<Self::WriteAccess<'_>> {
+        world.get_resource_mut::<R>()
+    }
+    unsafe fn fetch_read_only(world: UnsafeWorldCell<'_>) -> Option<Self::ReadOnlyAccess<'_>> {
+        world.get_resource::<R>()
+    }
+}
+
+// Allow the user to get `Res` access to a resource as well.
+// But getting `ResMut` isn't supported attow.
+impl<R: Resource> ResourceBundle for Res<'_, R> {
+    type WriteAccess<'a> = Mut<'a, R>;
+    type ReadOnlyAccess<'a> = Res<'a, R>;
+    unsafe fn fetch_write_access(world: UnsafeWorldCell<'_>) -> Option<Self::WriteAccess<'_>> {
+        world.get_resource_mut::<R>()
+    }
+    unsafe fn fetch_read_only(world: UnsafeWorldCell<'_>) -> Option<Self::ReadOnlyAccess<'_>> {
+        world.get_resource_ref::<R>()
+    }
+}
+
+impl<R: Resource> AccessConflictTracker for Res<'_, R> {
+    fn merge_with(other: &mut BundleAccessTable) {
+        other.insert_checked(TypeId::of::<R>());
+    }
+}
+
+impl<R: Resource> AccessConflictTracker for R {
+    fn merge_with(other: &mut BundleAccessTable) {
+        other.insert_checked(TypeId::of::<R>());
+    }
+}
+
+macro_rules! impl_conflict_tracker {
+    ($($tracker:ident),*) => {
+        impl <$($tracker: AccessConflictTracker),*> AccessConflictTracker for ($($tracker,)*) {
+            fn contains_conflicting_access() -> bool {
+                let mut tmp_table = BundleAccessTable::new_empty_unconflicted();
+                Self::merge_with(&mut tmp_table);
+                tmp_table.is_conflicted()
+            }
+
+            fn merge_with(other: &mut BundleAccessTable) {
+                $($tracker::merge_with(other));*
+            }
+        }
+    };
+}
+
+macro_rules! impl_resource_bundle {
+    ($($bundle:ident),*) => {
+        impl<$($bundle: ResourceBundle + AccessConflictTracker),*> ResourceBundle for ($($bundle,)*) {
+            type WriteAccess<'a> = ($($bundle::WriteAccess<'a>,)*);
+            type ReadOnlyAccess<'a> = ($($bundle::ReadOnlyAccess<'a>,)*);
+            unsafe fn fetch_write_access(world: UnsafeWorldCell<'_>) -> Option<Self::WriteAccess<'_>> {
+                Some(($($bundle::fetch_write_access(world)?,)*))
+            }
+            unsafe fn fetch_read_only(world: UnsafeWorldCell<'_>) -> Option<Self::ReadOnlyAccess<'_>> {
+                Some(($($bundle::fetch_read_only(world)?,)*))
+            }
+            fn contains_access_conflicts() -> bool {
+                <Self as AccessConflictTracker>::contains_conflicting_access()
+            }
+        }
+    };
+}
+
+all_tuples!(impl_resource_bundle, 1, 15, B);
+all_tuples!(impl_conflict_tracker, 1, 15, T);

crates/bevy_ecs/src/world/mod.rs

@@ -26,6 +26,7 @@ use crate::{
     event::{Event, EventId, Events, SendBatchIds},
     query::{DebugCheckedUnwrap, QueryData, QueryEntityError, QueryFilter, QueryState},
     removal_detection::RemovedComponentEvents,
+    resource_bundle::ResourceBundle,
     schedule::{Schedule, ScheduleLabel, Schedules},
     storage::{ResourceData, Storages},
     system::{Res, Resource},
@@ -1377,6 +1378,91 @@ impl World {
         unsafe { self.as_unsafe_world_cell().get_resource_mut() }
     }
 
+    /// Gets mutable access to multiple resources at once.
+    ///
+    /// Return `None` if one of the resources couldn't be fetched from the [`World`].
+    ///
+    /// # Panics
+    /// This method will panic if there are access conflicts within provided resource bundle.
+    /// For example, for any resources R, T, F:
+    /// `
+    /// world.get_resources_mut::<(R, R, T)>(); // This will panic!  There cannot be two mutable references to the same resource!
+    /// world.get_resources_mut::<(R, T, F)>(); // This is ok!
+    /// `
+    ///
+    /// # Examples
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    ///
+    /// #[derive(Resource)]
+    /// struct Num(isize);
+    ///
+    /// #[derive(Resource)]
+    /// struct BigNum(i128);
+    ///
+    /// let mut world = World::new();
+    /// world.insert_resource(Num(100));
+    /// world.insert_resource(BigNum(100_000));
+    ///
+    /// let (mut num, mut big_num) = world.get_resources_mut::<(Num, BigNum)>().unwrap();
+    /// num.0 *= 2;
+    /// big_num.0 *= 2;
+    ///
+    /// assert_eq!(world.resource::<Num>().0, 200);
+    /// assert_eq!(world.resource::<BigNum>().0, 200_000);
+    /// ```
+    #[inline]
+    pub fn get_resources_mut<B: ResourceBundle>(&mut self) -> Option<B::WriteAccess<'_>> {
+        assert!(!B::contains_access_conflicts(), "Found access conflicts in resource bundle. 
+            Make sure that if there is a mutable reference to some type R, it is the only reference to R in the bundle.");
+        // SAFETY: We have a mutable access to the world + we checked that there are no access conflicts within the bundle
+        unsafe { self.as_unsafe_world_cell().get_resources_mut::<B>() }
+    }
+
+    /// Gets read-only access to the resources in the bundle.
+    ///
+    /// Return `None` if one of the resources couldn't be fetched from the [`World`].
+    ///
+    /// # Examples
+    /// ```
+    /// use bevy_ecs::prelude::*;
+    ///
+    /// #[derive(Resource)]
+    /// struct Num(isize);
+    ///
+    /// #[derive(Resource)]
+    /// struct BigNum(i128);
+    ///
+    /// let mut world = World::new();
+    /// world.insert_resource(Num(100));
+    /// world.insert_resource(BigNum(100_000));
+    ///
+    /// let (Num(num), BigNum(big_num)) = world.get_resources::<(Num, BigNum)>().unwrap();
+    ///
+    /// assert_eq!(*num, 100);
+    /// assert_eq!(*big_num, 100_000);
+    /// ```
+    #[inline]
+    pub fn get_resources<B: ResourceBundle>(&self) -> Option<B::ReadOnlyAccess<'_>> {
+        // SAFETY: We have a shared reference to this `World`, so there aren't any other valid mutable references to the `World`'s resources.
+        unsafe { self.as_unsafe_world_cell_readonly().get_resources::<B>() }
+    }
+
+    /// Gets access to a bundle of resources at once, without checking for access conflicts within the bundle.
+    /// Similar to [`World::get_resources_mut`] but this will not check for access conflicts.
+    ///
+    /// # Safety
+    /// The caller must ensure that there are no access conflicts within the [`ResourceBundle`]. For example: (for any resources R, T, F)
+    /// (R, R)      -    *Not Allowed!*  Two mutable references can't exist at the same time!
+    /// (R, T, F)  -     *Allowed!*  No access conflicts
+    #[inline]
+    pub unsafe fn get_resources_mut_unchecked<B: ResourceBundle>(
+        &mut self,
+    ) -> Option<B::WriteAccess<'_>> {
+        self.as_unsafe_world_cell()
+            .get_resources_mut_unchecked::<B>()
+    }
+
     /// Gets a mutable reference to the resource of type `T` if it exists,
     /// otherwise inserts the resource using the result of calling `func`.
     #[inline]

crates/bevy_ecs/src/world/unsafe_world_cell.rs

@@ -13,6 +13,7 @@ use crate::{
     entity::{Entities, Entity, EntityLocation},
     prelude::Component,
     removal_detection::RemovedComponentEvents,
+    resource_bundle::ResourceBundle,
     storage::{Column, ComponentSparseSet, Storages},
     system::{Res, Resource},
 };
@@ -590,6 +591,60 @@ impl<'w> UnsafeWorldCell<'w> {
             .get(component_id)?
             .get_with_ticks()
     }
+
+    /// Gets mutable access to multiple resources at once.
+    ///
+    /// Return `None` if one of the resources couldn't be fetched from the [`World`].
+    ///
+    /// # Panics
+    /// This method will panic if there are access conflicts within provided resource bundle.
+    /// For example, for any resources R, T, F:
+    /// `
+    /// world.get_resources_mut::<(R, R, T)>(); // This will panic!  There cannot be two mutable references to the same resource!
+    /// world.get_resources_mut::<(R, T, F)>(); // This is ok!
+    /// `
+    /// See [`World::get_resources`] for examples.
+    ///
+    /// # Safety
+    /// It is the caller's responsibility to make sure that
+    /// - This [`UnsafeWorldCell`] has permission to access all of the resources.
+    /// - There are no other references to any of the resources.
+    #[inline]
+    pub unsafe fn get_resources_mut<B: ResourceBundle>(self) -> Option<B::WriteAccess<'w>> {
+        assert!(!B::contains_access_conflicts(), "Found access conflicts in resource bundle. 
+            Make sure that if there is a mutable reference to some type R, it is the only reference to R in the bundle.");
+        // SAFETY: The safety of this function as described in the "# Safety" section.
+        unsafe { B::fetch_write_access(self) }
+    }
+
+    /// Gets read-only access to the resources in the bundle.
+    ///
+    /// Return `None` if one of the resources couldn't be fetched from the [`World`].
+    /// See [`World::get_resources`] for examples.
+    /// # Safety
+    /// It is the caller's responsibility to make sure that
+    /// - This [`UnsafeWorldCell`] has permission to access all of the resources.
+    /// - There are no other mutable references to any of the resources.
+    #[inline]
+    pub unsafe fn get_resources<B: ResourceBundle>(self) -> Option<B::ReadOnlyAccess<'w>> {
+        // SAFETY: The safety of this function as described in the "# Safety" section.
+        unsafe { B::fetch_read_only(self) }
+    }
+
+    /// Gets access to a bundle of resources at once, without checking for access conflicts within the bundle.
+    /// Similar to [`World::get_resources_mut`] but this will not check for access conflicts.
+    ///
+    /// # Safety
+    /// The caller must ensure that there are no access conflicts within the [`ResourceBundle`]. For example: (for any resources R, T, F)
+    /// (R, R)      -    *Not Allowed!*  Two mutable references can't exist at the same time!
+    /// (R, T, F)  -     *Allowed!*  No access conflicts
+    #[inline]
+    pub unsafe fn get_resources_mut_unchecked<B: ResourceBundle>(
+        self,
+    ) -> Option<B::WriteAccess<'w>> {
+        // SAFETY: The safety of this function as described in the "# Safety" section.
+        unsafe { B::fetch_write_access(self) }
+    }
 }
 
 impl Debug for UnsafeWorldCell<'_> {

crates/bevy_utils/src/lib.rs

@@ -348,6 +348,10 @@ pub type EntityHashSet<T> = hashbrown::HashSet<T, EntityHash>;
 /// Iteration order only depends on the order of insertions and deletions.
 pub type TypeIdMap<V> = hashbrown::HashMap<TypeId, V, NoOpTypeIdHash>;
 
+/// A specialized hashset type with Key of [`TypeId`]
+/// Iteration order only depends on the order of insertions and deletions.
+pub type TypeIdSet = hashbrown::HashSet<TypeId, NoOpTypeIdHash>;
+
 /// [`BuildHasher`] for [`TypeId`]s.
 #[derive(Default)]
 pub struct NoOpTypeIdHash;