use bevy_utils::prelude::DebugName;
use crate::{
batching::BatchingStrategy,
change_detection::Tick,
entity::{Entity, EntityEquivalent, EntitySet, UniqueEntityArray},
query::{
DebugCheckedUnwrap, NopWorldQuery, QueryCombinationIter, QueryData, QueryEntityError,
QueryFilter, QueryIter, QueryManyIter, QueryManyUniqueIter, QueryParIter, QueryParManyIter,
QueryParManyUniqueIter, QuerySingleError, QueryState, ROQueryItem, ReadOnlyQueryData,
},
world::unsafe_world_cell::UnsafeWorldCell,
};
use core::{
marker::PhantomData,
mem::MaybeUninit,
ops::{Deref, DerefMut},
};
pub struct Query<'world, 'state, D: QueryData, F: QueryFilter = ()> {
world: UnsafeWorldCell<'world>,
state: &'state QueryState<D, F>,
last_run: Tick,
this_run: Tick,
}
impl<D: ReadOnlyQueryData, F: QueryFilter> Clone for Query<'_, '_, D, F> {
fn clone(&self) -> Self {
*self
}
}
impl<D: ReadOnlyQueryData, F: QueryFilter> Copy for Query<'_, '_, D, F> {}
impl<D: QueryData, F: QueryFilter> core::fmt::Debug for Query<'_, '_, D, F> {
fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
f.debug_struct("Query")
.field("matched_entities", &self.iter().count())
.field("state", &self.state)
.field("last_run", &self.last_run)
.field("this_run", &self.this_run)
.field("world", &self.world)
.finish()
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
#[inline]
pub(crate) unsafe fn new(
world: UnsafeWorldCell<'w>,
state: &'s QueryState<D, F>,
last_run: Tick,
this_run: Tick,
) -> Self {
Self {
world,
state,
last_run,
this_run,
}
}
pub fn as_readonly(&self) -> Query<'_, 's, D::ReadOnly, F> {
unsafe { self.reborrow_unsafe() }.into_readonly()
}
fn as_nop(&self) -> Query<'_, 's, NopWorldQuery<D>, F> {
let new_state = self.state.as_nop();
unsafe { Query::new(self.world, new_state, self.last_run, self.this_run) }
}
pub fn into_readonly(self) -> Query<'w, 's, D::ReadOnly, F> {
let new_state = self.state.as_readonly();
unsafe { Query::new(self.world, new_state, self.last_run, self.this_run) }
}
pub fn reborrow(&mut self) -> Query<'_, 's, D, F> {
unsafe { self.reborrow_unsafe() }
}
pub unsafe fn reborrow_unsafe(&self) -> Query<'_, 's, D, F> {
unsafe { self.copy_unsafe() }
}
unsafe fn copy_unsafe(&self) -> Query<'w, 's, D, F> {
unsafe { Query::new(self.world, self.state, self.last_run, self.this_run) }
}
#[inline]
pub fn iter(&self) -> QueryIter<'_, 's, D::ReadOnly, F> {
self.as_readonly().into_iter()
}
#[inline]
pub fn iter_mut(&mut self) -> QueryIter<'_, 's, D, F> {
self.reborrow().into_iter()
}
#[inline]
pub fn iter_combinations<const K: usize>(
&self,
) -> QueryCombinationIter<'_, 's, D::ReadOnly, F, K> {
self.as_readonly().iter_combinations_inner()
}
#[inline]
pub fn iter_combinations_mut<const K: usize>(
&mut self,
) -> QueryCombinationIter<'_, 's, D, F, K> {
self.reborrow().iter_combinations_inner()
}
#[inline]
pub fn iter_combinations_inner<const K: usize>(self) -> QueryCombinationIter<'w, 's, D, F, K> {
unsafe { QueryCombinationIter::new(self.world, self.state, self.last_run, self.this_run) }
}
#[inline]
pub fn iter_many<EntityList: IntoIterator<Item: EntityEquivalent>>(
&self,
entities: EntityList,
) -> QueryManyIter<'_, 's, D::ReadOnly, F, EntityList::IntoIter> {
self.as_readonly().iter_many_inner(entities)
}
#[inline]
pub fn iter_many_mut<EntityList: IntoIterator<Item: EntityEquivalent>>(
&mut self,
entities: EntityList,
) -> QueryManyIter<'_, 's, D, F, EntityList::IntoIter> {
self.reborrow().iter_many_inner(entities)
}
#[inline]
pub fn iter_many_inner<EntityList: IntoIterator<Item: EntityEquivalent>>(
self,
entities: EntityList,
) -> QueryManyIter<'w, 's, D, F, EntityList::IntoIter> {
unsafe {
QueryManyIter::new(
self.world,
self.state,
entities,
self.last_run,
self.this_run,
)
}
}
#[inline]
pub fn iter_many_unique<EntityList: EntitySet>(
&self,
entities: EntityList,
) -> QueryManyUniqueIter<'_, 's, D::ReadOnly, F, EntityList::IntoIter> {
self.as_readonly().iter_many_unique_inner(entities)
}
#[inline]
pub fn iter_many_unique_mut<EntityList: EntitySet>(
&mut self,
entities: EntityList,
) -> QueryManyUniqueIter<'_, 's, D, F, EntityList::IntoIter> {
self.reborrow().iter_many_unique_inner(entities)
}
#[inline]
pub fn iter_many_unique_inner<EntityList: EntitySet>(
self,
entities: EntityList,
) -> QueryManyUniqueIter<'w, 's, D, F, EntityList::IntoIter> {
unsafe {
QueryManyUniqueIter::new(
self.world,
self.state,
entities,
self.last_run,
self.this_run,
)
}
}
#[inline]
pub unsafe fn iter_unsafe(&self) -> QueryIter<'_, 's, D, F> {
unsafe { self.reborrow_unsafe() }.into_iter()
}
#[inline]
pub unsafe fn iter_combinations_unsafe<const K: usize>(
&self,
) -> QueryCombinationIter<'_, 's, D, F, K> {
unsafe { self.reborrow_unsafe() }.iter_combinations_inner()
}
pub unsafe fn iter_many_unsafe<EntityList: IntoIterator<Item: EntityEquivalent>>(
&self,
entities: EntityList,
) -> QueryManyIter<'_, 's, D, F, EntityList::IntoIter> {
unsafe { self.reborrow_unsafe() }.iter_many_inner(entities)
}
pub unsafe fn iter_many_unique_unsafe<EntityList: EntitySet>(
&self,
entities: EntityList,
) -> QueryManyUniqueIter<'_, 's, D, F, EntityList::IntoIter> {
unsafe { self.reborrow_unsafe() }.iter_many_unique_inner(entities)
}
#[inline]
pub fn par_iter(&self) -> QueryParIter<'_, 's, D::ReadOnly, F> {
self.as_readonly().par_iter_inner()
}
#[inline]
pub fn par_iter_mut(&mut self) -> QueryParIter<'_, 's, D, F> {
self.reborrow().par_iter_inner()
}
#[inline]
pub fn par_iter_inner(self) -> QueryParIter<'w, 's, D, F> {
QueryParIter {
world: self.world,
state: self.state,
last_run: self.last_run,
this_run: self.this_run,
batching_strategy: BatchingStrategy::new(),
}
}
#[inline]
pub fn par_iter_many<EntityList: IntoIterator<Item: EntityEquivalent>>(
&self,
entities: EntityList,
) -> QueryParManyIter<'_, 's, D::ReadOnly, F, EntityList::Item> {
QueryParManyIter {
world: self.world,
state: self.state.as_readonly(),
entity_list: entities.into_iter().collect(),
last_run: self.last_run,
this_run: self.this_run,
batching_strategy: BatchingStrategy::new(),
}
}
#[inline]
pub fn par_iter_many_unique<EntityList: EntitySet<Item: Sync>>(
&self,
entities: EntityList,
) -> QueryParManyUniqueIter<'_, 's, D::ReadOnly, F, EntityList::Item> {
QueryParManyUniqueIter {
world: self.world,
state: self.state.as_readonly(),
entity_list: entities.into_iter().collect(),
last_run: self.last_run,
this_run: self.this_run,
batching_strategy: BatchingStrategy::new(),
}
}
#[inline]
pub fn par_iter_many_unique_mut<EntityList: EntitySet<Item: Sync>>(
&mut self,
entities: EntityList,
) -> QueryParManyUniqueIter<'_, 's, D, F, EntityList::Item> {
QueryParManyUniqueIter {
world: self.world,
state: self.state,
entity_list: entities.into_iter().collect(),
last_run: self.last_run,
this_run: self.this_run,
batching_strategy: BatchingStrategy::new(),
}
}
#[inline]
pub fn get(&self, entity: Entity) -> Result<ROQueryItem<'_, 's, D>, QueryEntityError> {
self.as_readonly().get_inner(entity)
}
#[inline]
pub fn get_many<const N: usize>(
&self,
entities: [Entity; N],
) -> Result<[ROQueryItem<'_, 's, D>; N], QueryEntityError> {
self.as_readonly().get_many_inner(entities)
}
#[inline]
pub fn get_many_unique<const N: usize>(
&self,
entities: UniqueEntityArray<N>,
) -> Result<[ROQueryItem<'_, 's, D>; N], QueryEntityError> {
self.as_readonly().get_many_unique_inner(entities)
}
#[inline]
pub fn get_mut(&mut self, entity: Entity) -> Result<D::Item<'_, 's>, QueryEntityError> {
self.reborrow().get_inner(entity)
}
#[inline]
pub fn get_inner(self, entity: Entity) -> Result<D::Item<'w, 's>, QueryEntityError> {
unsafe {
let location = self.world.entities().get_spawned(entity)?;
if !self
.state
.matched_archetypes
.contains(location.archetype_id.index())
{
return Err(QueryEntityError::QueryDoesNotMatch(
entity,
location.archetype_id,
));
}
let archetype = self
.world
.archetypes()
.get(location.archetype_id)
.debug_checked_unwrap();
let mut fetch = D::init_fetch(
self.world,
&self.state.fetch_state,
self.last_run,
self.this_run,
);
let mut filter = F::init_fetch(
self.world,
&self.state.filter_state,
self.last_run,
self.this_run,
);
let table = self
.world
.storages()
.tables
.get(location.table_id)
.debug_checked_unwrap();
D::set_archetype(&mut fetch, &self.state.fetch_state, archetype, table);
F::set_archetype(&mut filter, &self.state.filter_state, archetype, table);
if F::filter_fetch(
&self.state.filter_state,
&mut filter,
entity,
location.table_row,
) && let Some(item) = D::fetch(
&self.state.fetch_state,
&mut fetch,
entity,
location.table_row,
) {
Ok(item)
} else {
Err(QueryEntityError::QueryDoesNotMatch(
entity,
location.archetype_id,
))
}
}
}
#[inline]
pub fn get_many_mut<const N: usize>(
&mut self,
entities: [Entity; N],
) -> Result<[D::Item<'_, 's>; N], QueryEntityError> {
self.reborrow().get_many_mut_inner(entities)
}
#[inline]
pub fn get_many_unique_mut<const N: usize>(
&mut self,
entities: UniqueEntityArray<N>,
) -> Result<[D::Item<'_, 's>; N], QueryEntityError> {
self.reborrow().get_many_unique_inner(entities)
}
#[inline]
pub fn get_many_mut_inner<const N: usize>(
self,
entities: [Entity; N],
) -> Result<[D::Item<'w, 's>; N], QueryEntityError> {
for i in 0..N {
for j in 0..i {
if entities[i] == entities[j] {
return Err(QueryEntityError::AliasedMutability(entities[i]));
}
}
}
unsafe { self.get_many_impl(entities) }
}
#[inline]
pub fn get_many_inner<const N: usize>(
self,
entities: [Entity; N],
) -> Result<[D::Item<'w, 's>; N], QueryEntityError>
where
D: ReadOnlyQueryData,
{
unsafe { self.get_many_impl(entities) }
}
#[inline]
pub fn get_many_unique_inner<const N: usize>(
self,
entities: UniqueEntityArray<N>,
) -> Result<[D::Item<'w, 's>; N], QueryEntityError> {
unsafe { self.get_many_impl(entities.into_inner()) }
}
unsafe fn get_many_impl<const N: usize>(
self,
entities: [Entity; N],
) -> Result<[D::Item<'w, 's>; N], QueryEntityError> {
let mut values = [(); N].map(|_| MaybeUninit::uninit());
for (value, entity) in core::iter::zip(&mut values, entities) {
let item = unsafe { self.copy_unsafe() }.get_inner(entity)?;
*value = MaybeUninit::new(item);
}
Ok(values.map(|x| unsafe { x.assume_init() }))
}
#[inline]
pub unsafe fn get_unchecked(
&self,
entity: Entity,
) -> Result<D::Item<'_, 's>, QueryEntityError> {
unsafe { self.reborrow_unsafe() }.get_inner(entity)
}
#[inline]
pub fn single(&self) -> Result<ROQueryItem<'_, 's, D>, QuerySingleError> {
self.as_readonly().single_inner()
}
#[inline]
pub fn single_mut(&mut self) -> Result<D::Item<'_, 's>, QuerySingleError> {
self.reborrow().single_inner()
}
#[inline]
pub fn single_inner(self) -> Result<D::Item<'w, 's>, QuerySingleError> {
let mut query = self.into_iter();
let first = query.next();
let extra = query.next().is_some();
match (first, extra) {
(Some(r), false) => Ok(r),
(None, _) => Err(QuerySingleError::NoEntities(DebugName::type_name::<Self>())),
(Some(_), _) => Err(QuerySingleError::MultipleEntities(DebugName::type_name::<
Self,
>())),
}
}
#[inline]
pub fn is_empty(&self) -> bool {
if D::IS_ARCHETYPAL {
self.as_nop().iter().next().is_none()
} else {
self.iter().next().is_none()
}
}
#[inline]
pub fn contains(&self, entity: Entity) -> bool {
if D::IS_ARCHETYPAL {
self.as_nop().get(entity).is_ok()
} else {
self.get(entity).is_ok()
}
}
pub fn count(&self) -> usize {
if !D::IS_ARCHETYPAL {
self.into_iter().count()
} else if !F::IS_ARCHETYPAL {
self.as_nop().into_iter().count()
} else {
self.as_nop().into_iter().size_hint().0
}
}
#[track_caller]
pub fn transmute_lens<NewD: QueryData>(&mut self) -> QueryLens<'_, NewD> {
self.transmute_lens_filtered::<NewD, ()>()
}
#[track_caller]
pub fn transmute_lens_inner<NewD: QueryData>(self) -> QueryLens<'w, NewD> {
self.transmute_lens_filtered_inner::<NewD, ()>()
}
#[track_caller]
pub fn transmute_lens_filtered<NewD: QueryData, NewF: QueryFilter>(
&mut self,
) -> QueryLens<'_, NewD, NewF> {
self.reborrow().transmute_lens_filtered_inner()
}
#[track_caller]
pub fn transmute_lens_filtered_inner<NewD: QueryData, NewF: QueryFilter>(
self,
) -> QueryLens<'w, NewD, NewF> {
let state = self.state.transmute_filtered::<NewD, NewF>(self.world);
QueryLens {
world: self.world,
state,
last_run: self.last_run,
this_run: self.this_run,
}
}
pub fn as_query_lens(&mut self) -> QueryLens<'_, D> {
self.transmute_lens()
}
pub fn into_query_lens(self) -> QueryLens<'w, D> {
self.transmute_lens_inner()
}
pub fn join<'a, OtherD: QueryData, NewD: QueryData>(
&'a mut self,
other: &'a mut Query<OtherD>,
) -> QueryLens<'a, NewD> {
self.join_filtered(other)
}
pub fn join_inner<OtherD: QueryData, NewD: QueryData>(
self,
other: Query<'w, '_, OtherD>,
) -> QueryLens<'w, NewD> {
self.join_filtered_inner(other)
}
pub fn join_filtered<
'a,
OtherD: QueryData,
OtherF: QueryFilter,
NewD: QueryData,
NewF: QueryFilter,
>(
&'a mut self,
other: &'a mut Query<OtherD, OtherF>,
) -> QueryLens<'a, NewD, NewF> {
self.reborrow().join_filtered_inner(other.reborrow())
}
pub fn join_filtered_inner<
OtherD: QueryData,
OtherF: QueryFilter,
NewD: QueryData,
NewF: QueryFilter,
>(
self,
other: Query<'w, '_, OtherD, OtherF>,
) -> QueryLens<'w, NewD, NewF> {
let state = self
.state
.join_filtered::<OtherD, OtherF, NewD, NewF>(self.world, other.state);
QueryLens {
world: self.world,
state,
last_run: self.last_run,
this_run: self.this_run,
}
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for Query<'w, 's, D, F> {
type Item = D::Item<'w, 's>;
type IntoIter = QueryIter<'w, 's, D, F>;
fn into_iter(self) -> Self::IntoIter {
unsafe { QueryIter::new(self.world, self.state, self.last_run, self.this_run) }
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'w Query<'_, 's, D, F> {
type Item = ROQueryItem<'w, 's, D>;
type IntoIter = QueryIter<'w, 's, D::ReadOnly, F>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'w mut Query<'_, 's, D, F> {
type Item = D::Item<'w, 's>;
type IntoIter = QueryIter<'w, 's, D, F>;
fn into_iter(self) -> Self::IntoIter {
self.iter_mut()
}
}
impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Query<'w, 's, D, F> {
#[inline]
pub fn iter_inner(&self) -> QueryIter<'w, 's, D::ReadOnly, F> {
(*self).into_iter()
}
}
pub struct QueryLens<'w, Q: QueryData, F: QueryFilter = ()> {
world: UnsafeWorldCell<'w>,
state: QueryState<Q, F>,
last_run: Tick,
this_run: Tick,
}
impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> {
pub fn query(&mut self) -> Query<'_, '_, Q, F> {
Query {
world: self.world,
state: &self.state,
last_run: self.last_run,
this_run: self.this_run,
}
}
}
impl<'w, Q: ReadOnlyQueryData, F: QueryFilter> QueryLens<'w, Q, F> {
pub fn query_inner(&self) -> Query<'w, '_, Q, F> {
Query {
world: self.world,
state: &self.state,
last_run: self.last_run,
this_run: self.this_run,
}
}
}
impl<'w, 's, Q: QueryData, F: QueryFilter> From<&'s mut QueryLens<'w, Q, F>>
for Query<'s, 's, Q, F>
{
fn from(value: &'s mut QueryLens<'w, Q, F>) -> Query<'s, 's, Q, F> {
value.query()
}
}
impl<'w, 'q, Q: QueryData, F: QueryFilter> From<&'q mut Query<'w, '_, Q, F>>
for QueryLens<'q, Q, F>
{
fn from(value: &'q mut Query<'w, '_, Q, F>) -> QueryLens<'q, Q, F> {
value.transmute_lens_filtered()
}
}
pub struct Single<'w, 's, D: QueryData, F: QueryFilter = ()> {
pub(crate) item: D::Item<'w, 's>,
pub(crate) _filter: PhantomData<F>,
}
impl<'w, 's, D: QueryData, F: QueryFilter> Deref for Single<'w, 's, D, F> {
type Target = D::Item<'w, 's>;
fn deref(&self) -> &Self::Target {
&self.item
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> DerefMut for Single<'w, 's, D, F> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.item
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> Single<'w, 's, D, F> {
pub fn into_inner(self) -> D::Item<'w, 's> {
self.item
}
}
pub struct Populated<'w, 's, D: QueryData, F: QueryFilter = ()>(pub(crate) Query<'w, 's, D, F>);
impl<'w, 's, D: QueryData, F: QueryFilter> Deref for Populated<'w, 's, D, F> {
type Target = Query<'w, 's, D, F>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<D: QueryData, F: QueryFilter> DerefMut for Populated<'_, '_, D, F> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> Populated<'w, 's, D, F> {
pub fn into_inner(self) -> Query<'w, 's, D, F> {
self.0
}
}
impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for Populated<'w, 's, D, F> {
type Item = <Query<'w, 's, D, F> as IntoIterator>::Item;
type IntoIter = <Query<'w, 's, D, F> as IntoIterator>::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
impl<'a, 'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'a Populated<'w, 's, D, F> {
type Item = <&'a Query<'w, 's, D, F> as IntoIterator>::Item;
type IntoIter = <&'a Query<'w, 's, D, F> as IntoIterator>::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.deref().into_iter()
}
}
impl<'a, 'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'a mut Populated<'w, 's, D, F> {
type Item = <&'a mut Query<'w, 's, D, F> as IntoIterator>::Item;
type IntoIter = <&'a mut Query<'w, 's, D, F> as IntoIterator>::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.deref_mut().into_iter()
}
}
#[cfg(test)]
mod tests {
use crate::{prelude::*, query::QueryEntityError};
use alloc::vec::Vec;
#[test]
fn get_many_uniqueness() {
let mut world = World::new();
let entities: Vec<Entity> = (0..10).map(|_| world.spawn_empty().id()).collect();
let mut query_state = world.query::<Entity>();
assert!(query_state
.query_mut(&mut world)
.get_many_mut_inner::<10>(entities.clone().try_into().unwrap())
.is_ok());
assert_eq!(
query_state
.query_mut(&mut world)
.get_many_mut_inner([entities[0], entities[0]])
.unwrap_err(),
QueryEntityError::AliasedMutability(entities[0])
);
assert_eq!(
query_state
.query_mut(&mut world)
.get_many_mut_inner([entities[0], entities[1], entities[0]])
.unwrap_err(),
QueryEntityError::AliasedMutability(entities[0])
);
assert_eq!(
query_state
.query_mut(&mut world)
.get_many_mut_inner([entities[9], entities[9]])
.unwrap_err(),
QueryEntityError::AliasedMutability(entities[9])
);
}
}