use crate::{
component::{ComponentId, Components, StorageType},
query::FilteredAccess,
world::{FromWorld, World},
};
use bevy_ecs_macros::{Component, Resource};
use smallvec::SmallVec;
#[cfg(feature = "bevy_reflect")]
use {
crate::reflect::{ReflectComponent, ReflectResource},
bevy_reflect::std_traits::ReflectDefault,
bevy_reflect::Reflect,
};
#[derive(Component, Clone, Debug, Default)]
#[cfg_attr(
feature = "bevy_reflect",
derive(Reflect),
reflect(Component),
reflect(Debug, Clone, Default)
)]
pub struct Disabled;
#[derive(Resource, Debug)]
#[cfg_attr(
feature = "bevy_reflect",
derive(bevy_reflect::Reflect),
reflect(Resource)
)]
pub struct DefaultQueryFilters {
disabling: SmallVec<[ComponentId; 4]>,
}
impl FromWorld for DefaultQueryFilters {
fn from_world(world: &mut World) -> Self {
let mut filters = DefaultQueryFilters::empty();
let disabled_component_id = world.register_component::<Disabled>();
filters.register_disabling_component(disabled_component_id);
filters
}
}
impl DefaultQueryFilters {
#[must_use]
pub fn empty() -> Self {
DefaultQueryFilters {
disabling: SmallVec::new(),
}
}
pub fn register_disabling_component(&mut self, component_id: ComponentId) {
if !self.disabling.contains(&component_id) {
self.disabling.push(component_id);
}
}
pub fn disabling_ids(&self) -> impl Iterator<Item = ComponentId> {
self.disabling.iter().copied()
}
pub(super) fn modify_access(&self, component_access: &mut FilteredAccess) {
for component_id in self.disabling_ids() {
if !component_access.contains(component_id) {
component_access.and_without(component_id);
}
}
}
pub(super) fn is_dense(&self, components: &Components) -> bool {
self.disabling_ids().all(|component_id| {
components
.get_info(component_id)
.is_some_and(|info| info.storage_type() == StorageType::Table)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
prelude::{EntityMut, EntityRef, World},
query::{Has, With},
};
use alloc::{vec, vec::Vec};
#[test]
fn filters_modify_access() {
let mut filters = DefaultQueryFilters::empty();
filters.register_disabling_component(ComponentId::new(1));
let mut component_access = FilteredAccess::default();
component_access.access_mut().add_read(ComponentId::new(2));
let mut applied_access = component_access.clone();
filters.modify_access(&mut applied_access);
assert_eq!(0, applied_access.with_filters().count());
assert_eq!(
vec![ComponentId::new(1)],
applied_access.without_filters().collect::<Vec<_>>()
);
component_access.and_with(ComponentId::new(4));
let mut applied_access = component_access.clone();
filters.modify_access(&mut applied_access);
assert_eq!(
vec![ComponentId::new(4)],
applied_access.with_filters().collect::<Vec<_>>()
);
assert_eq!(
vec![ComponentId::new(1)],
applied_access.without_filters().collect::<Vec<_>>()
);
let copy = component_access.clone();
component_access.and_with(ComponentId::new(1));
let mut applied_access = component_access.clone();
filters.modify_access(&mut applied_access);
assert_eq!(
vec![ComponentId::new(1), ComponentId::new(4)],
applied_access.with_filters().collect::<Vec<_>>()
);
assert_eq!(0, applied_access.without_filters().count());
component_access = copy.clone();
component_access
.access_mut()
.add_archetypal(ComponentId::new(1));
let mut applied_access = component_access.clone();
filters.modify_access(&mut applied_access);
assert_eq!(
vec![ComponentId::new(4)],
applied_access.with_filters().collect::<Vec<_>>()
);
assert_eq!(0, applied_access.without_filters().count());
}
#[derive(Component)]
struct CustomDisabled;
#[derive(Component)]
struct Dummy;
#[test]
fn multiple_disabling_components() {
let mut world = World::new();
world.register_disabling_component::<CustomDisabled>();
world.spawn(Dummy);
world.spawn_batch((0..2).map(|_| (Dummy, Disabled)));
world.spawn_batch((0..4).map(|_| (Dummy, CustomDisabled)));
world.spawn_batch((0..8).map(|_| (Dummy, Disabled, CustomDisabled)));
let mut query = world.query::<&Dummy>();
assert_eq!(1, query.iter(&world).count());
let mut query = world.query_filtered::<EntityRef, With<Dummy>>();
assert_eq!(1, query.iter(&world).count());
let mut query = world.query_filtered::<EntityMut, With<Dummy>>();
assert_eq!(1, query.iter(&world).count());
let mut query = world.query_filtered::<&Dummy, With<Disabled>>();
assert_eq!(2, query.iter(&world).count());
let mut query = world.query_filtered::<Has<Disabled>, With<Dummy>>();
assert_eq!(3, query.iter(&world).count());
let mut query = world.query_filtered::<&Dummy, With<CustomDisabled>>();
assert_eq!(4, query.iter(&world).count());
let mut query = world.query_filtered::<Has<CustomDisabled>, With<Dummy>>();
assert_eq!(5, query.iter(&world).count());
let mut query = world.query_filtered::<&Dummy, (With<Disabled>, With<CustomDisabled>)>();
assert_eq!(8, query.iter(&world).count());
let mut query = world.query_filtered::<(Has<Disabled>, Has<CustomDisabled>), With<Dummy>>();
assert_eq!(15, query.iter(&world).count());
let mut query = world.query_filtered::<Option<&Disabled>, With<Dummy>>();
assert_eq!(1, query.iter(&world).count());
}
}