Allow despawning while running queries (#37)

* Manual rustfmt

* Allow debug-querying list of components

* Some simplification of the common entity removal code path

* Allow despawning while running queries

Author: oli-obk

src/query.rs

@@ -95,13 +95,18 @@ impl<C: 'static> SparseSetGetter for &mut C {
 macro_rules! impl_query {
     ($($T:ident),*) => {
         impl<A: SparseSetGetter + Always, $($T: SparseSetGetter,)* Z> Query<(A, $($T,)*)> for Z
-where
-    Z: FnMut(Entity, A::Short<'_>, $($T::Short<'_>,)*) + FnMut(Entity, A, $($T,)*),{
+        where
+            Z: FnMut(Entity, A::Short<'_>, $($T::Short<'_>,)*),
+            Z: FnMut(Entity, A, $($T,)*),
+        {
             #[track_caller]
             fn get_components(world: &World, mut f: Z) {
                 #[allow(non_snake_case)]
                 if let (Some(mut a), $(Some(mut $T),)*) = (A::get_set(world), $($T::get_set(world),)*) {
                     for (entity, a) in A::iter(&mut a) {
+                        if world.is_despawning(entity) {
+                            continue;
+                        }
                         $(let Some($T) = $T::get_entity(&mut $T, entity) else { continue };)*
                         f(entity, a, $($T,)*);
 

src/scheduler.rs

@@ -55,6 +55,7 @@ impl<'a, RES> Scheduler<'a, RES> {
             let mut sys = sys.take().unwrap();
             drop(guard);
             sys(world, res);
+            world.flush_despawned();
             let mut guard = self.systems.borrow_mut();
             let Some((_, entry)) = guard.get_mut(i) else {
                 break;

src/sparse_set.rs

@@ -35,17 +35,16 @@ impl<C> SparseSet<C> {
         let idx = self.sparse.remove(&entity)?;
         let last = self.dense.len() - 1;
 
-        if idx != last {
-            self.dense.swap(idx, last);
-            let entity = *self.ids.last().unwrap();
-            self.ids.swap(idx, last);
+        if idx == last {
+            self.ids.pop();
+            self.dense.pop()
+        } else {
+            self.ids.swap_remove(idx);
 
-            let _prev = self.sparse.insert(entity, idx);
+            let _prev = self.sparse.insert(self.ids[idx], idx);
             debug_assert_eq!(_prev, Some(last));
+            Some(self.dense.swap_remove(idx))
         }
-
-        self.ids.pop();
-        self.dense.pop()
     }
 
     pub fn get(&self, entity: Entity) -> Option<&C> {
@@ -74,20 +73,16 @@ impl<C> SparseSet<C> {
 }
 
 trait Set: Any {
-    #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-    fn remove(&mut self, entity: Entity) -> Option<&'static str>;
+    fn debug(&mut self, entity: Entity) -> Option<&'static str>;
 
-    #[cfg(not(any(debug_assertions, feature = "track_dead_entities")))]
     fn remove(&mut self, entity: Entity);
 }
 
 impl<C: Any> Set for SparseSet<C> {
-    #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-    fn remove(&mut self, entity: Entity) -> Option<&'static str> {
-        self.remove(entity).map(|_| type_name::<C>())
+    fn debug(&mut self, entity: Entity) -> Option<&'static str> {
+        self.get(entity).map(|_| type_name::<C>())
     }
 
-    #[cfg(not(any(debug_assertions, feature = "track_dead_entities")))]
     fn remove(&mut self, entity: Entity) {
         self.remove(entity);
     }
@@ -99,12 +94,6 @@ pub struct SparseSets {
     sets: FrozenVec<Box<RefCell<dyn Set>>>,
 }
 
-#[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-pub(crate) type RemoveType = String;
-
-#[cfg(not(any(debug_assertions, feature = "track_dead_entities")))]
-pub(crate) type RemoveType = ();
-
 impl SparseSets {
     pub fn insert<C: Any>(&self, entity: Entity, component: C) {
         let component = Box::new(RefCell::new(SparseSet::new(entity, component)));
@@ -117,7 +106,7 @@ impl SparseSets {
     }
 
     #[track_caller]
-    pub fn remove(&self, entity: Entity) -> RemoveType {
+    pub fn debug(&self, entity: Entity) -> String {
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
         let mut component = String::new();
 
@@ -129,7 +118,7 @@ impl SparseSets {
             };
 
             #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-            if let Some(c) = guard.remove(entity) {
+            if let Some(c) = guard.debug(entity) {
                 component.push_str(c);
                 component.push_str(", ");
             }
@@ -140,6 +129,19 @@ impl SparseSets {
         component
     }
 
+    #[track_caller]
+    pub fn remove(&self, entity: Entity) {
+        for set in self.sets.iter() {
+            let Ok(mut guard) = set.try_borrow_mut() else {
+                panic!(
+                    "Tried to access component mutably, but it is already being read or written to",
+                )
+            };
+
+            guard.remove(entity);
+        }
+    }
+
     #[track_caller]
     pub fn get<C: 'static>(&self) -> Option<Ref<SparseSet<C>>> {
         let i = *self.set_access.borrow().get(&TypeId::of::<C>())?;

src/world.rs

@@ -13,7 +13,7 @@ use std::{
 use crate::{
     query::Query,
     scheduler::{Scheduler, SysId},
-    sparse_set::{RemoveType, SparseSets},
+    sparse_set::SparseSets,
 };
 
 /// An opaque id for an entity.
@@ -56,7 +56,13 @@ impl Default for EntityCounter {
 pub struct World {
     entities: EntityCounter,
     #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-    dead_entities: RefCell<BTreeMap<Entity, (&'static Location<'static>, String)>>,
+    dead_entity_locations: RefCell<BTreeMap<Entity, &'static Location<'static>>>,
+    #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
+    dead_entity_components: RefCell<BTreeMap<Entity, String>>,
+    /// Entities that are being despawned.
+    /// As we cannot mutate the sparse sets to remove all components while queries are running, we instead do
+    /// so whenever a system finishes
+    despawning: RefCell<Vec<Entity>>,
     pub(crate) sparse_sets: SparseSets,
     scheduler: Scheduler<'static, ()>,
 }
@@ -65,7 +71,8 @@ impl World {
     #[track_caller]
     fn insert<C: Any>(&self, entity: Entity, component: C) {
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-        if let Some((loc, components)) = self.dead_entities.borrow().get(&entity) {
+        if let Some(components) = self.dead_entity_components.borrow().get(&entity) {
+            let loc = self.dead_entity_locations.borrow()[&entity];
             panic!(
                 "Attaching `{}` to despawned entity (despawned at {loc}).Components at despawn time: {components}",
                 type_name::<C>(),
@@ -94,11 +101,29 @@ impl World {
     /// Destroy an entity and all its components. Future attempts to use this entity in any way will panic.
     #[track_caller]
     pub fn despawn(&self, entity: Entity) {
-        let _detach_info = self.detach_all(entity);
+        self.despawning.borrow_mut().push(entity);
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-        self.dead_entities
+        self.dead_entity_locations
             .borrow_mut()
-            .insert(entity, (Location::caller(), _detach_info));
+            .insert(entity, Location::caller());
+    }
+
+    /// Ensure all despawned entities actually got removed
+    #[track_caller]
+    pub fn flush_despawned(&self) {
+        for entity in self.despawning.borrow_mut().drain(..) {
+            #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
+            let components = self.debug_components(entity);
+            self.detach_all(entity);
+            #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
+            self.dead_entity_components
+                .borrow_mut()
+                .insert(entity, components);
+        }
+    }
+
+    pub(crate) fn is_despawning(&self, entity: Entity) -> bool {
+        self.despawning.borrow().contains(&entity)
     }
 
     /// Attach multiple components to an entity at once.
@@ -111,7 +136,8 @@ impl World {
     #[track_caller]
     pub fn detach<C: 'static>(&self, entity: Entity) -> Option<C> {
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-        if let Some((loc, components)) = self.dead_entities.borrow().get(&entity) {
+        if let Some(components) = self.dead_entity_components.borrow().get(&entity) {
+            let loc = self.dead_entity_locations.borrow()[&entity];
             panic!(
                 "Detaching `{}` from despawned entity (despawned at {loc})\nComponents at despawn time: {components}",
                 type_name::<C>(),
@@ -124,16 +150,24 @@ impl World {
     /// Detach all components from an entity and drop them.
     /// If you want to extract specific components, call [Self::detach] first.
     #[track_caller]
-    pub fn detach_all(&self, entity: Entity) -> RemoveType {
+    pub fn detach_all(&self, entity: Entity) {
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-        if let Some((loc, components)) = self.dead_entities.borrow().get(&entity) {
+        if let Some(components) = self.dead_entity_components.borrow().get(&entity) {
+            let loc = self.dead_entity_locations.borrow()[&entity];
             panic!(
                 "Detaching all components from despawned entity (despawned at {loc})\nComponents at despawn time: {components}"
             );
         }
         self.sparse_sets.remove(entity)
     }
 
+    /// Detach all components from an entity and drop them.
+    /// If you want to extract specific components, call [Self::detach] first.
+    #[track_caller]
+    pub fn debug_components(&self, entity: Entity) -> String {
+        self.sparse_sets.debug(entity)
+    }
+
     /// Detach all components of a specific type from all entities and drop them.
     ///
     /// This method removes all components of type `C` from every entity in the world.
@@ -169,7 +203,8 @@ impl World {
     #[track_caller]
     pub fn get<C: 'static>(&self, entity: Entity) -> Option<Ref<C>> {
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-        if let Some((loc, components)) = self.dead_entities.borrow().get(&entity) {
+        if let Some(components) = self.dead_entity_components.borrow().get(&entity) {
+            let loc = self.dead_entity_locations.borrow()[&entity];
             panic!(
                 "Getting `{}` from despawned entity (despawned at {loc})\nComponents at despawn time: {components}",
                 type_name::<C>(),
@@ -187,7 +222,8 @@ impl World {
     #[track_caller]
     pub fn get_mut<C: 'static>(&self, entity: Entity) -> Option<RefMut<C>> {
         #[cfg(any(debug_assertions, feature = "track_dead_entities"))]
-        if let Some((loc, components)) = self.dead_entities.borrow().get(&entity) {
+        if let Some(components) = self.dead_entity_components.borrow().get(&entity) {
+            let loc = self.dead_entity_locations.borrow()[&entity];
             panic!(
                 "Getting `{}` from despawned entity (despawned at {loc})\nComponents at despawn time: {components}",
                 type_name::<C>(),

tests/entity.rs

@@ -57,6 +57,16 @@ fn spawn_related_in_query() {
 fn attach_to_despawned() {
     let world = World::default();
 
+    let entity = world.spawn((1_u32,));
+    world.despawn(entity);
+    world.flush_despawned();
+    world.attach(entity, (String::new(),));
+}
+
+#[test]
+fn attach_to_despawned_no_flush() {
+    let world = World::default();
+
     let entity = world.spawn((1_u32,));
     world.despawn(entity);
     world.attach(entity, (String::new(),));

tests/system.rs

@@ -43,6 +43,27 @@ fn despawn() {
     assert_eq!(&results[..], &[(10, Some("foo"))]);
 }
 
+#[test]
+fn despawn_in_query() {
+    let world = World::default();
+
+    let id = world.spawn((1_u32,));
+    world.spawn((10_u32, "foo"));
+    world.despawn(id);
+
+    let mut results = vec![];
+    world.query(|entity, i: &u32, s: Option<&&'static str>| {
+        results.push((*i, s.map(|s| *s)));
+        world.despawn(entity);
+    });
+    assert_eq!(&results[..], &[(10, Some("foo"))]);
+    let mut results = vec![];
+    world.query(|_, i: &u32, s: Option<&&'static str>| {
+        results.push((*i, s.map(|s| *s)));
+    });
+    assert_eq!(&results[..], &[]);
+}
+
 #[test]
 fn get() {
     let world = World::default();

tests/ui/panic_location.rs

@@ -8,6 +8,7 @@ fn main() {
     let id = world.spawn((1_u32,));
     world.spawn((10_u32, "foo"));
     world.despawn(id);
+    world.flush_despawned();
 
     world.get::<u32>(id);
 }

tests/ui/panic_location.run.stderr

@@ -1,5 +1,5 @@
 
-thread 'main' panicked at tests/ui/panic_location.rs:12:11:
+thread 'main' panicked at tests/ui/panic_location.rs:13:11:
 Getting `u32` from despawned entity (despawned at tests/ui/panic_location.rs:10:11)
 Components at despawn time: u32, 
 note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace