Add methods for iterating resources

crates/bevy_ecs/src/storage/resource.rs

@@ -50,10 +50,7 @@ impl<const SEND: bool> ResourceData<SEND> {
     /// If `SEND` is false, this will panic if called from a different thread than the one it was inserted from.
     #[inline]
     fn validate_access(&self) {
-        if SEND {
-            return;
-        }
-        if self.origin_thread_id != Some(std::thread::current().id()) {
+        if !self.can_access() {
             // Panic in tests, as testing for aborting is nearly impossible
             panic!(
                 "Attempted to access or drop non-send resource {} from thread {:?} on a thread {:?}. This is not allowed. Aborting.",
@@ -64,6 +61,18 @@ impl<const SEND: bool> ResourceData<SEND> {
         }
     }
 
+    /// Returns whether the current thread can access a resource or not:
+    ///
+    /// - For `!Send` resources, this will return `true` if the current thread is the same as the one it was inserted from;
+    /// - For `Send` resources, this will always return `true`.
+    #[inline]
+    pub fn can_access(&self) -> bool {
+        if SEND {
+            return true;
+        }
+        self.origin_thread_id == Some(std::thread::current().id())
+    }
+
     /// Returns true if the resource is populated.
     #[inline]
     pub fn is_present(&self) -> bool {

crates/bevy_ecs/src/world/mod.rs

@@ -27,7 +27,7 @@ use crate::{
     query::{DebugCheckedUnwrap, QueryData, QueryEntityError, QueryFilter, QueryState},
     removal_detection::RemovedComponentEvents,
     schedule::{Schedule, ScheduleLabel, Schedules},
-    storage::{ResourceData, Storages},
+    storage::{ResourceData, Resources, Storages},
     system::{Res, Resource},
     world::error::TryRunScheduleError,
 };
@@ -2040,6 +2040,79 @@ impl World {
         self.storages.resources.clear();
         self.storages.non_send_resources.clear();
     }
+
+    /// Iterates over all resources in the world
+    pub fn iter_resources(&self) -> impl Iterator<Item = (&ComponentInfo, Ptr<'_>)> {
+        self.iter_resources_internal(&self.storages.resources)
+    }
+
+    /// Mutably iterates over all resources in the world
+    pub fn iter_resources_mut(&mut self) -> impl Iterator<Item = (&ComponentInfo, MutUntyped<'_>)> {
+        // SAFETY: `&mut self` ensures that all accessed data is unaliased
+        unsafe { self.iter_resources_mut_internal(&self.storages.resources) }
+    }
+
+    /// Iterates over all `!Send` resources in the world that are accessible from the current thread.
+    ///
+    /// Resources that were inserted from a different thread are skipped.
+    pub fn iter_non_send(&self) -> impl Iterator<Item = (&ComponentInfo, Ptr<'_>)> {
+        self.iter_resources_internal(&self.storages.non_send_resources)
+    }
+
+    /// Mutably iterates over all `!Send` resources in the world that are accessible from the current thread.
+    ///
+    /// Resources that were inserted from a different thread are skipped.
+    pub fn iter_non_send_mut(&mut self) -> impl Iterator<Item = (&ComponentInfo, MutUntyped<'_>)> {
+        // SAFETY: `&mut self` ensures that all accessed data is unaliased
+        unsafe { self.iter_resources_mut_internal(&self.storages.non_send_resources) }
+    }
+
+    fn iter_resources_internal<'w, const SEND: bool>(
+        &'w self,
+        resources: &'w Resources<SEND>,
+    ) -> impl Iterator<Item = (&ComponentInfo, Ptr<'_>)> {
+        resources.iter().filter(|(_, data)| data.can_access()).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)
+        })
+    }
+
+    /// # Safety
+    /// - Caller must have exclusive access to the world
+    unsafe fn iter_resources_mut_internal<'w, const SEND: bool>(
+        &'w self,
+        resources: &'w Resources<SEND>,
+    ) -> impl Iterator<Item = (&ComponentInfo, MutUntyped<'_>)> {
+        resources.iter().filter(|(_, data)| data.can_access()).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()
+                )
+            });
+
+            let ticks = TicksMut::from_tick_cells(ticks, self.last_change_tick(), self.read_change_tick());
+
+            let mut_untyped = MutUntyped {
+                value: ptr.assert_unique(),
+                ticks,
+            };
+
+            (component_info, mut_untyped)
+        })
+    }
 }
 
 impl World {