From 92a0bafc24d6cef43554bf90309b923841c867c9 Mon Sep 17 00:00:00 2001
From: mike <ramirezmike2@gmail.com>
Date: Mon, 17 Oct 2022 13:47:00 +0000
Subject: [PATCH] Add iter_entities to World #6228 (#6242)

# Objective

- Add a way to iterate over all entities from &World

## Solution

- Added a function `iter_entities` on World which returns an iterator of `Entity` derived from the entities in the `World`'s `archetypes`

---

## Changelog

- Added a function `iter_entities` on World, allowing iterating over all entities in contexts where you only have read-only access to the World.
---
 crates/bevy_ecs/src/world/mod.rs       | 81 +++++++++++++++++++++++++-
 crates/bevy_scene/src/dynamic_scene.rs |  7 +--
 2 files changed, 81 insertions(+), 7 deletions(-)

diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs
index 128c5e8dcf094..3e45e8c100fea 100644
--- a/crates/bevy_ecs/src/world/mod.rs
+++ b/crates/bevy_ecs/src/world/mod.rs
@@ -320,6 +320,16 @@ impl World {
         Some(EntityRef::new(self, entity, location))
     }
 
+    /// Returns an [`Entity`] iterator of current entities.
+    ///
+    /// This is useful in contexts where you only have read-only access to the [`World`].
+    #[inline]
+    pub fn iter_entities(&self) -> impl Iterator<Item = Entity> + '_ {
+        self.archetypes
+            .iter()
+            .flat_map(|archetype| archetype.entities().iter().copied())
+    }
+
     /// Retrieves an [`EntityMut`] that exposes read and write operations for the given `entity`.
     /// Returns [`None`] if the `entity` does not exist. Use [`World::entity_mut`] if you don't want
     /// to unwrap the [`EntityMut`] yourself.
@@ -1634,7 +1644,7 @@ mod tests {
         system::Resource,
     };
     use bevy_ecs_macros::Component;
-    use bevy_utils::HashSet;
+    use bevy_utils::{HashMap, HashSet};
     use std::{
         any::TypeId,
         panic,
@@ -1907,4 +1917,73 @@ mod tests {
             [Some(baz_id)].into()
         );
     }
+
+    #[test]
+    fn iterate_entities() {
+        let mut world = World::new();
+        let mut entity_counters = HashMap::new();
+
+        let iterate_and_count_entities = |world: &World, entity_counters: &mut HashMap<_, _>| {
+            entity_counters.clear();
+            for entity in world.iter_entities() {
+                let counter = entity_counters.entry(entity).or_insert(0);
+                *counter += 1;
+            }
+        };
+
+        // Adding one entity and validating iteration
+        let ent0 = world.spawn((Foo, Bar, Baz)).id();
+
+        iterate_and_count_entities(&world, &mut entity_counters);
+        assert_eq!(entity_counters[&ent0], 1);
+        assert_eq!(entity_counters.len(), 1);
+
+        // Spawning three more entities and then validating iteration
+        let ent1 = world.spawn((Foo, Bar)).id();
+        let ent2 = world.spawn((Bar, Baz)).id();
+        let ent3 = world.spawn((Foo, Baz)).id();
+
+        iterate_and_count_entities(&world, &mut entity_counters);
+
+        assert_eq!(entity_counters[&ent0], 1);
+        assert_eq!(entity_counters[&ent1], 1);
+        assert_eq!(entity_counters[&ent2], 1);
+        assert_eq!(entity_counters[&ent3], 1);
+        assert_eq!(entity_counters.len(), 4);
+
+        // Despawning first entity and then validating the iteration
+        assert!(world.despawn(ent0));
+
+        iterate_and_count_entities(&world, &mut entity_counters);
+
+        assert_eq!(entity_counters[&ent1], 1);
+        assert_eq!(entity_counters[&ent2], 1);
+        assert_eq!(entity_counters[&ent3], 1);
+        assert_eq!(entity_counters.len(), 3);
+
+        // Spawning three more entities, despawning three and then validating the iteration
+        let ent4 = world.spawn(Foo).id();
+        let ent5 = world.spawn(Bar).id();
+        let ent6 = world.spawn(Baz).id();
+
+        assert!(world.despawn(ent2));
+        assert!(world.despawn(ent3));
+        assert!(world.despawn(ent4));
+
+        iterate_and_count_entities(&world, &mut entity_counters);
+
+        assert_eq!(entity_counters[&ent1], 1);
+        assert_eq!(entity_counters[&ent5], 1);
+        assert_eq!(entity_counters[&ent6], 1);
+        assert_eq!(entity_counters.len(), 3);
+
+        // Despawning remaining entities and then validating the iteration
+        assert!(world.despawn(ent1));
+        assert!(world.despawn(ent5));
+        assert!(world.despawn(ent6));
+
+        iterate_and_count_entities(&world, &mut entity_counters);
+
+        assert_eq!(entity_counters.len(), 0);
+    }
 }
diff --git a/crates/bevy_scene/src/dynamic_scene.rs b/crates/bevy_scene/src/dynamic_scene.rs
index 16b7a3169cf7f..1cd92c284bfd2 100644
--- a/crates/bevy_scene/src/dynamic_scene.rs
+++ b/crates/bevy_scene/src/dynamic_scene.rs
@@ -42,12 +42,7 @@ impl DynamicScene {
         let mut builder =
             DynamicSceneBuilder::from_world_with_type_registry(world, type_registry.clone());
 
-        builder.extract_entities(
-            world
-                .archetypes()
-                .iter()
-                .flat_map(|archetype| archetype.entities().iter().copied()),
-        );
+        builder.extract_entities(world.iter_entities());
 
         builder.build()
     }