Add Disabled marker

crates/bevy_ecs/src/lib.rs

@@ -45,7 +45,9 @@ pub mod prelude {
         component::Component,
         entity::{Entity, EntityMapper},
         event::{Event, EventReader, EventWriter, Events},
-        query::{Added, AnyOf, Changed, Has, Or, QueryBuilder, QueryState, With, Without},
+        query::{
+            Added, AnyOf, Changed, Disabled, Has, Or, QueryBuilder, QueryState, With, Without,
+        },
         removal_detection::RemovedComponents,
         schedule::{
             apply_deferred, apply_state_transition, common_conditions::*, Condition,
@@ -69,7 +71,7 @@ mod tests {
         change_detection::Ref,
         component::{Component, ComponentId},
         entity::Entity,
-        query::{Added, Changed, FilteredAccess, QueryFilter, With, Without},
+        query::{Added, Changed, Disabled, FilteredAccess, QueryFilter, With, Without},
         system::Resource,
         world::{EntityRef, Mut, World},
     };
@@ -1379,13 +1381,15 @@ mod tests {
     #[test]
     fn filtered_query_access() {
         let mut world = World::new();
+        let disabled_id = world.init_component::<Disabled>();
         let query = world.query_filtered::<&mut A, Changed<B>>();
 
         let mut expected = FilteredAccess::<ComponentId>::default();
         let a_id = world.components.get_id(TypeId::of::<A>()).unwrap();
         let b_id = world.components.get_id(TypeId::of::<B>()).unwrap();
         expected.add_write(a_id);
         expected.add_read(b_id);
+        expected.and_without(disabled_id);
         assert!(
             query.component_access.eq(&expected),
             "ComponentId access from query fetch and query filter should be combined"

crates/bevy_ecs/src/query/mod.rs

@@ -21,6 +21,12 @@ pub use par_iter::*;
 pub use state::*;
 pub use world_query::*;
 
+use crate as bevy_ecs;
+
+/// A marker component that excludes an entity from queries that don't specifically request them.
+#[derive(bevy_ecs_macros::Component)]
+pub struct Disabled;
+
 /// A debug checked version of [`Option::unwrap_unchecked`]. Will panic in
 /// debug modes if unwrapping a `None` or `Err` value in debug mode, but is
 /// equivalent to `Option::unwrap_unchecked` or `Result::unwrap_unchecked`

crates/bevy_ecs/src/query/state.rs

@@ -2,7 +2,7 @@ use crate::{
     archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
     component::{ComponentId, Tick},
     entity::Entity,
-    prelude::FromWorld,
+    prelude::{Disabled, FromWorld},
     query::{
         Access, BatchingStrategy, DebugCheckedUnwrap, FilteredAccess, QueryCombinationIter,
         QueryIter, QueryParIter,
@@ -146,6 +146,16 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     }
 }
 
+fn contains_disabled(
+    component_access: &FilteredAccess<ComponentId>,
+    disabled_id: ComponentId,
+) -> bool {
+    component_access.access().has_read(disabled_id)
+        || component_access.filter_sets.iter().any(|f| {
+            f.with.contains(disabled_id.index()) || f.without.contains(disabled_id.index())
+        })
+}
+
 impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
     /// Creates a new [`QueryState`] from a given [`World`] and inherits the result of `world.id()`.
     pub fn new(world: &mut World) -> Self {
@@ -190,6 +200,12 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
         // properly considered in a global "cross-query" context (both within systems and across systems).
         component_access.extend(&filter_component_access);
 
+        let disabled_id = world.init_component::<Disabled>();
+        let allow_disabled = contains_disabled(&component_access, disabled_id);
+        if !allow_disabled {
+            component_access.and_without(disabled_id);
+        }
+
         Self {
             world_id: world.id(),
             archetype_generation: ArchetypeGeneration::initial(),
@@ -214,13 +230,20 @@ impl<D: QueryData, F: QueryFilter> QueryState<D, F> {
         let filter_state = F::init_state(builder.world_mut());
         D::set_access(&mut fetch_state, builder.access());
 
+        let mut component_access = builder.access().clone();
+        let disabled_id = builder.world_mut().init_component::<Disabled>();
+        let allow_disabled = contains_disabled(&component_access, disabled_id);
+        if !allow_disabled {
+            component_access.and_without(disabled_id);
+        }
+
         let mut state = Self {
             world_id: builder.world().id(),
             archetype_generation: ArchetypeGeneration::initial(),
             matched_storage_ids: Vec::new(),
             fetch_state,
             filter_state,
-            component_access: builder.access().clone(),
+            component_access,
             matched_tables: Default::default(),
             matched_archetypes: Default::default(),
             #[cfg(feature = "trace")]
@@ -1627,7 +1650,7 @@ impl<D: QueryData, F: QueryFilter> From<QueryBuilder<'_, D, F>> for QueryState<D
 mod tests {
     use crate as bevy_ecs;
     use crate::world::FilteredEntityRef;
-    use crate::{component::Component, prelude::*, query::QueryEntityError};
+    use crate::{prelude::*, query::QueryEntityError};
 
     #[test]
     fn get_many_unchecked_manual_uniqueness() {
@@ -1999,4 +2022,47 @@ mod tests {
         let query_2 = QueryState::<&B, Without<C>>::new(&mut world);
         let _: QueryState<Entity, Changed<C>> = query_1.join_filtered(&world, &query_2);
     }
+
+    #[test]
+    fn does_exclude_disabled() {
+        let mut world = World::new();
+        let entity_a = world.spawn(A(0)).id();
+        let _ = world.spawn((A(1), Disabled)).id();
+        let entity_c = world.spawn(A(1)).id();
+        world.disable(entity_c);
+
+        let mut query = QueryState::<Entity>::new(&mut world);
+        assert_eq!(entity_a, query.single(&world));
+
+        let mut query = QueryState::<Entity, Without<Disabled>>::new(&mut world);
+        assert_eq!(entity_a, query.single(&world));
+    }
+
+    #[test]
+    fn can_request_disabled() {
+        let mut world = World::new();
+        let _ = world.spawn(A(0)).id();
+        let _ = world.spawn((A(1), Disabled)).id();
+        let entity_c = world.spawn(A(1)).id();
+        world.disable(entity_c);
+
+        let mut query = QueryState::<Entity, With<Disabled>>::new(&mut world);
+        assert_eq!(2, query.iter(&world).count());
+
+        let mut query = QueryState::<(Entity, &Disabled)>::new(&mut world);
+        assert_eq!(2, query.iter(&world).count());
+
+        let mut query = QueryState::<(Entity, &mut Disabled)>::new(&mut world);
+        assert_eq!(2, query.iter(&world).count());
+
+        let mut query = QueryState::<(Entity, Ref<Disabled>)>::new(&mut world);
+        assert_eq!(2, query.iter(&world).count());
+
+        let mut query = QueryState::<(Entity, Option<&Disabled>)>::new(&mut world);
+        assert_eq!(3, query.iter(&world).count());
+
+        // TODO: This case is not handled correctly yet, since Has<T> does not register component access
+        // let mut query = QueryState::<(Entity, Has<Disabled>)>::new(&mut world);
+        // assert_eq!(3, query.iter(&world).count());
+    }
 }

crates/bevy_ecs/src/system/commands/mod.rs

@@ -933,6 +933,16 @@ impl EntityCommands<'_> {
         self.add(despawn);
     }
 
+    /// Enable this entity, removing the [`bevy_ecs::prelude::Disabled`] marker if it's present
+    pub fn enable(&mut self) {
+        self.add(enable);
+    }
+
+    /// Disable this entity, adding the [`bevy_ecs::prelude::Disabled`] marker to the entity
+    pub fn disable(&mut self) {
+        self.add(disable);
+    }
+
     /// Pushes an [`EntityCommand`] to the queue, which will get executed for the current [`Entity`].
     ///
     /// # Examples
@@ -1084,6 +1094,14 @@ fn despawn(entity: Entity, world: &mut World) {
     world.despawn(entity);
 }
 
+fn enable(entity: Entity, world: &mut World) {
+    world.enable(entity);
+}
+
+fn disable(entity: Entity, world: &mut World) {
+    world.disable(entity);
+}
+
 /// An [`EntityCommand`] that adds the components in a [`Bundle`] to an entity.
 fn insert<T: Bundle>(bundle: T) -> impl EntityCommand {
     move |entity: Entity, world: &mut World| {

crates/bevy_ecs/src/world/entity_ref.rs

@@ -4,7 +4,7 @@ use crate::{
     change_detection::MutUntyped,
     component::{Component, ComponentId, ComponentTicks, Components, StorageType},
     entity::{Entities, Entity, EntityLocation},
-    query::{Access, DebugCheckedUnwrap},
+    query::{Access, DebugCheckedUnwrap, Disabled},
     removal_detection::RemovedComponentEvents,
     storage::Storages,
     world::{Mut, World},
@@ -1136,6 +1136,8 @@ impl<'w> EntityWorldMut<'w> {
     ///
     /// See [`EntityCommands::retain`](crate::system::EntityCommands::retain) for more details.
     pub fn retain<T: Bundle>(&mut self) -> &mut Self {
+        let disabled_id = self.world.init_component::<Disabled>();
+
         let archetypes = &mut self.world.archetypes;
         let storages = &mut self.world.storages;
         let components = &mut self.world.components;
@@ -1148,7 +1150,7 @@ impl<'w> EntityWorldMut<'w> {
 
         let to_remove = &old_archetype
             .components()
-            .filter(|c| !retained_bundle_info.components().contains(c))
+            .filter(|c| !retained_bundle_info.components().contains(c) && *c != disabled_id)
             .collect::<Vec<_>>();
         let remove_bundle = self.world.bundles.init_dynamic_info(components, to_remove);
 
@@ -1269,6 +1271,16 @@ impl<'w> EntityWorldMut<'w> {
         self.entity
     }
 
+    /// Enables the current entity, removing the [`Disabled`] marker.
+    pub fn enable(&mut self) {
+        self.remove::<Disabled>();
+    }
+
+    /// Disables the current entity, adding the [`Disabled`] marker.
+    pub fn disable(&mut self) {
+        self.insert(Disabled);
+    }
+
     /// Gets read-only access to the world that the current entity belongs to.
     #[inline]
     pub fn world(&self) -> &World {
@@ -2583,6 +2595,33 @@ mod tests {
         assert_eq!(world.entity(ent).archetype().components().next(), None);
     }
 
+    // Test that calling retain retains [`Disabled`].
+    #[test]
+    fn retain_disabled() {
+        #[derive(Component)]
+        struct Marker<const N: usize>;
+
+        let mut world = World::new();
+        let ent = world
+            .spawn((Marker::<1>, Marker::<2>, Marker::<3>, Disabled))
+            .id();
+
+        let disabled_id = world.init_component::<Disabled>();
+        world.entity_mut(ent).retain::<Marker<1>>();
+        assert!(world
+            .entity(ent)
+            .archetype()
+            .components()
+            .any(|c| c == disabled_id));
+
+        world.entity_mut(ent).retain::<()>();
+        assert!(world
+            .entity(ent)
+            .archetype()
+            .components()
+            .any(|c| c == disabled_id))
+    }
+
     // Test removing some components with `retain`, including components not on the entity.
     #[test]
     fn retain_some_components() {

crates/bevy_ecs/src/world/mod.rs

@@ -939,6 +939,32 @@ impl World {
         }
     }
 
+    /// Enables the given `entity`, if it exists and was [`crate::prelude::Disabled`].
+    /// Returns `true` if the `entity` is now enabled and `false` if the `entity` does not exist.
+    #[inline]
+    pub fn enable(&mut self, entity: Entity) -> bool {
+        if let Some(mut entity) = self.get_entity_mut(entity) {
+            entity.enable();
+            true
+        } else {
+            warn!("error[B0003]: Could not enable entity {:?} because it doesn't exist in this World.", entity);
+            false
+        }
+    }
+
+    /// Disabled the given `entity`, if it exists and didn't have [`crate::prelude::Disabled`] yet.
+    /// Returns `true` if the `entity` is now disabled and `false` if the `entity` does not exist.
+    #[inline]
+    pub fn disable(&mut self, entity: Entity) -> bool {
+        if let Some(mut entity) = self.get_entity_mut(entity) {
+            entity.disable();
+            true
+        } else {
+            warn!("error[B0003]: Could not disable entity {:?} because it doesn't exist in this World.", entity);
+            false
+        }
+    }
+
     /// Clears the internal component tracker state.
     ///
     /// The world maintains some internal state about changed and removed components. This state