use core::{
marker::PhantomData,
ops::{Deref, DerefMut},
ptr::NonNull,
};
use crate::{archetype::Archetype, Component, Entity};
pub trait Query {
#[doc(hidden)]
type Fetch: for<'a> Fetch<'a>;
}
pub trait Fetch<'a>: Sized {
type Item;
fn access(archetype: &Archetype) -> Option<Access>;
fn borrow(archetype: &Archetype);
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self>;
fn release(archetype: &Archetype);
unsafe fn should_skip(&self) -> bool {
false
}
unsafe fn next(&mut self) -> Self::Item;
}
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub enum Access {
Iterate,
Read,
Write,
}
#[derive(Copy, Clone, Debug)]
pub struct EntityFetch(NonNull<u32>);
impl Query for Entity {
type Fetch = EntityFetch;
}
impl<'a> Fetch<'a> for EntityFetch {
type Item = Entity;
#[inline]
fn access(_archetype: &Archetype) -> Option<Access> {
Some(Access::Iterate)
}
#[inline]
fn borrow(_archetype: &Archetype) {}
#[inline]
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
Some(EntityFetch(NonNull::new_unchecked(
archetype.entities().as_ptr().add(offset),
)))
}
#[inline]
fn release(_archetype: &Archetype) {}
#[inline]
unsafe fn next(&mut self) -> Self::Item {
let id = self.0.as_ptr();
self.0 = NonNull::new_unchecked(id.add(1));
Entity::from_id(*id)
}
}
impl<'a, T: Component> Query for &'a T {
type Fetch = FetchRead<T>;
}
#[doc(hidden)]
pub struct FetchRead<T>(NonNull<T>);
impl<'a, T: Component> Fetch<'a> for FetchRead<T> {
type Item = &'a T;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Read)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow::<T>();
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
archetype
.get::<T>()
.map(|x| Self(NonNull::new_unchecked(x.as_ptr().add(offset))))
}
fn release(archetype: &Archetype) {
archetype.release::<T>();
}
#[inline]
unsafe fn next(&mut self) -> &'a T {
let x = self.0.as_ptr();
self.0 = NonNull::new_unchecked(x.add(1));
&*x
}
}
impl<'a, T: Component> Query for &'a mut T {
type Fetch = FetchMut<T>;
}
impl<T: Query> Query for Option<T> {
type Fetch = TryFetch<T::Fetch>;
}
pub struct Mut<'a, T: Component> {
value: &'a mut T,
mutated: &'a mut bool,
}
unsafe impl<T: Component> Send for Mut<'_, T> {}
unsafe impl<T: Component> Sync for Mut<'_, T> {}
impl<'a, T: Component> Deref for Mut<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
self.value
}
}
impl<'a, T: Component> DerefMut for Mut<'a, T> {
#[inline]
fn deref_mut(&mut self) -> &mut T {
*self.mutated = true;
self.value
}
}
impl<'a, T: Component + core::fmt::Debug> core::fmt::Debug for Mut<'a, T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
self.value.fmt(f)
}
}
impl<'a, T: Component> Query for Mut<'a, T> {
type Fetch = FetchMut<T>;
}
#[doc(hidden)]
pub struct FetchMut<T>(NonNull<T>, NonNull<bool>);
impl<'a, T: Component> Fetch<'a> for FetchMut<T> {
type Item = Mut<'a, T>;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Write)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow_mut::<T>();
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
archetype
.get_with_mutated::<T>()
.map(|(components, mutated)| {
Self(
NonNull::new_unchecked(components.as_ptr().add(offset)),
NonNull::new_unchecked(mutated.as_ptr().add(offset)),
)
})
}
fn release(archetype: &Archetype) {
archetype.release_mut::<T>();
}
#[inline]
unsafe fn next(&mut self) -> Mut<'a, T> {
let component = self.0.as_ptr();
let mutated = self.1.as_ptr();
self.0 = NonNull::new_unchecked(component.add(1));
self.1 = NonNull::new_unchecked(mutated.add(1));
Mut {
value: &mut *component,
mutated: &mut *mutated,
}
}
}
#[allow(missing_docs)]
pub struct Mutated<'a, T> {
value: &'a T,
}
impl<'a, T: Component> Deref for Mutated<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
self.value
}
}
impl<'a, T: Component> Query for Mutated<'a, T> {
type Fetch = FetchMutated<T>;
}
#[doc(hidden)]
pub struct FetchMutated<T>(NonNull<T>, NonNull<bool>);
impl<'a, T: Component> Fetch<'a> for FetchMutated<T> {
type Item = Mutated<'a, T>;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Read)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow::<T>();
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
archetype
.get_with_mutated::<T>()
.map(|(components, mutated)| {
Self(
NonNull::new_unchecked(components.as_ptr().add(offset)),
NonNull::new_unchecked(mutated.as_ptr().add(offset)),
)
})
}
fn release(archetype: &Archetype) {
archetype.release::<T>();
}
unsafe fn should_skip(&self) -> bool {
!*self.1.as_ref()
}
#[inline]
unsafe fn next(&mut self) -> Self::Item {
self.1 = NonNull::new_unchecked(self.1.as_ptr().add(1));
let value = self.0.as_ptr();
self.0 = NonNull::new_unchecked(value.add(1));
Mutated { value: &*value }
}
}
#[allow(missing_docs)]
pub struct Added<'a, T> {
value: &'a T,
}
impl<'a, T: Component> Deref for Added<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
self.value
}
}
impl<'a, T: Component> Query for Added<'a, T> {
type Fetch = FetchAdded<T>;
}
#[doc(hidden)]
pub struct FetchAdded<T>(NonNull<T>, NonNull<bool>);
impl<'a, T: Component> Fetch<'a> for FetchAdded<T> {
type Item = Added<'a, T>;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Read)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow::<T>();
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
archetype.get_with_added::<T>().map(|(components, added)| {
Self(
NonNull::new_unchecked(components.as_ptr().add(offset)),
NonNull::new_unchecked(added.as_ptr().add(offset)),
)
})
}
fn release(archetype: &Archetype) {
archetype.release::<T>();
}
unsafe fn should_skip(&self) -> bool {
!*self.1.as_ref()
}
#[inline]
unsafe fn next(&mut self) -> Self::Item {
self.1 = NonNull::new_unchecked(self.1.as_ptr().add(1));
let value = self.0.as_ptr();
self.0 = NonNull::new_unchecked(value.add(1));
Added { value: &*value }
}
}
#[allow(missing_docs)]
pub struct Changed<'a, T> {
value: &'a T,
}
impl<'a, T: Component> Deref for Changed<'a, T> {
type Target = T;
#[inline]
fn deref(&self) -> &T {
self.value
}
}
impl<'a, T: Component> Query for Changed<'a, T> {
type Fetch = FetchChanged<T>;
}
#[doc(hidden)]
pub struct FetchChanged<T>(NonNull<T>, NonNull<bool>, NonNull<bool>);
impl<'a, T: Component> Fetch<'a> for FetchChanged<T> {
type Item = Changed<'a, T>;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Read)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow::<T>();
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
archetype
.get_with_added_and_mutated::<T>()
.map(|(components, added, mutated)| {
Self(
NonNull::new_unchecked(components.as_ptr().add(offset)),
NonNull::new_unchecked(added.as_ptr().add(offset)),
NonNull::new_unchecked(mutated.as_ptr().add(offset)),
)
})
}
fn release(archetype: &Archetype) {
archetype.release::<T>();
}
unsafe fn should_skip(&self) -> bool {
!*self.1.as_ref() && !self.2.as_ref()
}
#[inline]
unsafe fn next(&mut self) -> Self::Item {
self.1 = NonNull::new_unchecked(self.1.as_ptr().add(1));
self.2 = NonNull::new_unchecked(self.2.as_ptr().add(1));
let value = self.0.as_ptr();
self.0 = NonNull::new_unchecked(value.add(1));
Changed { value: &*value }
}
}
#[doc(hidden)]
pub struct TryFetch<T>(Option<T>);
impl<'a, T: Fetch<'a>> Fetch<'a> for TryFetch<T> {
type Item = Option<T::Item>;
fn access(archetype: &Archetype) -> Option<Access> {
Some(T::access(archetype).unwrap_or(Access::Iterate))
}
fn borrow(archetype: &Archetype) {
T::borrow(archetype)
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
Some(Self(T::get(archetype, offset)))
}
fn release(archetype: &Archetype) {
T::release(archetype)
}
unsafe fn next(&mut self) -> Option<T::Item> {
Some(self.0.as_mut()?.next())
}
unsafe fn should_skip(&self) -> bool {
self.0.as_ref().map_or(false, |fetch| fetch.should_skip())
}
}
pub struct Without<T, Q>(PhantomData<(Q, fn(T))>);
impl<T: Component, Q: Query> Query for Without<T, Q> {
type Fetch = FetchWithout<T, Q::Fetch>;
}
#[doc(hidden)]
pub struct FetchWithout<T, F>(F, PhantomData<fn(T)>);
impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWithout<T, F> {
type Item = F::Item;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
None
} else {
F::access(archetype)
}
}
fn borrow(archetype: &Archetype) {
F::borrow(archetype)
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
if archetype.has::<T>() {
return None;
}
Some(Self(F::get(archetype, offset)?, PhantomData))
}
fn release(archetype: &Archetype) {
F::release(archetype)
}
unsafe fn next(&mut self) -> F::Item {
self.0.next()
}
unsafe fn should_skip(&self) -> bool {
self.0.should_skip()
}
}
pub struct With<T, Q>(PhantomData<(Q, fn(T))>);
impl<T: Component, Q: Query> Query for With<T, Q> {
type Fetch = FetchWith<T, Q::Fetch>;
}
#[doc(hidden)]
pub struct FetchWith<T, F>(F, PhantomData<fn(T)>);
impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWith<T, F> {
type Item = F::Item;
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
F::access(archetype)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
F::borrow(archetype)
}
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
if !archetype.has::<T>() {
return None;
}
Some(Self(F::get(archetype, offset)?, PhantomData))
}
fn release(archetype: &Archetype) {
F::release(archetype)
}
unsafe fn next(&mut self) -> F::Item {
self.0.next()
}
unsafe fn should_skip(&self) -> bool {
self.0.should_skip()
}
}
pub struct QueryBorrow<'w, Q: Query> {
archetypes: &'w [Archetype],
borrowed: bool,
_marker: PhantomData<Q>,
}
impl<'w, Q: Query> QueryBorrow<'w, Q> {
pub(crate) fn new(archetypes: &'w [Archetype]) -> Self {
Self {
archetypes,
borrowed: false,
_marker: PhantomData,
}
}
pub fn iter<'q>(&'q mut self) -> QueryIter<'q, 'w, Q> {
self.borrow();
QueryIter {
borrow: self,
archetype_index: 0,
iter: None,
}
}
pub fn iter_batched<'q>(&'q mut self, batch_size: u32) -> BatchedIter<'q, 'w, Q> {
self.borrow();
BatchedIter {
borrow: self,
archetype_index: 0,
batch_size,
batch: 0,
}
}
fn borrow(&mut self) {
if self.borrowed {
panic!(
"called QueryBorrow::iter twice on the same borrow; construct a new query instead"
);
}
for x in self.archetypes {
if Q::Fetch::access(x) >= Some(Access::Read) {
Q::Fetch::borrow(x);
}
}
self.borrowed = true;
}
pub fn with<T: Component>(self) -> QueryBorrow<'w, With<T, Q>> {
self.transform()
}
pub fn without<T: Component>(self) -> QueryBorrow<'w, Without<T, Q>> {
self.transform()
}
fn transform<R: Query>(mut self) -> QueryBorrow<'w, R> {
let x = QueryBorrow {
archetypes: self.archetypes,
borrowed: self.borrowed,
_marker: PhantomData,
};
self.borrowed = false;
x
}
}
unsafe impl<'w, Q: Query> Send for QueryBorrow<'w, Q> {}
unsafe impl<'w, Q: Query> Sync for QueryBorrow<'w, Q> {}
impl<'w, Q: Query> Drop for QueryBorrow<'w, Q> {
fn drop(&mut self) {
if self.borrowed {
for x in self.archetypes {
if Q::Fetch::access(x) >= Some(Access::Read) {
Q::Fetch::release(x);
}
}
}
}
}
impl<'q, 'w, Q: Query> IntoIterator for &'q mut QueryBorrow<'w, Q> {
type IntoIter = QueryIter<'q, 'w, Q>;
type Item = <Q::Fetch as Fetch<'q>>::Item;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub struct QueryIter<'q, 'w, Q: Query> {
borrow: &'q mut QueryBorrow<'w, Q>,
archetype_index: u32,
iter: Option<ChunkIter<Q>>,
}
unsafe impl<'q, 'w, Q: Query> Send for QueryIter<'q, 'w, Q> {}
unsafe impl<'q, 'w, Q: Query> Sync for QueryIter<'q, 'w, Q> {}
impl<'q, 'w, Q: Query> Iterator for QueryIter<'q, 'w, Q> {
type Item = <Q::Fetch as Fetch<'q>>::Item;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
loop {
match self.iter {
None => {
let archetype = self.borrow.archetypes.get(self.archetype_index as usize)?;
self.archetype_index += 1;
unsafe {
self.iter = Q::Fetch::get(archetype, 0).map(|fetch| ChunkIter {
fetch,
len: archetype.len(),
});
}
}
Some(ref mut iter) => match unsafe { iter.next() } {
None => {
self.iter = None;
continue;
}
Some(components) => {
return Some(components);
}
},
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let n = self.len();
(n, Some(n))
}
}
impl<'q, 'w, Q: Query> ExactSizeIterator for QueryIter<'q, 'w, Q> {
fn len(&self) -> usize {
self.borrow
.archetypes
.iter()
.filter(|&x| Q::Fetch::access(x).is_some())
.map(|x| x.len() as usize)
.sum()
}
}
struct ChunkIter<Q: Query> {
fetch: Q::Fetch,
len: u32,
}
impl<Q: Query> ChunkIter<Q> {
unsafe fn next<'a>(&mut self) -> Option<<Q::Fetch as Fetch<'a>>::Item> {
loop {
if self.len == 0 {
return None;
}
self.len -= 1;
if self.fetch.should_skip() {
let _ = self.fetch.next();
continue;
}
break Some(self.fetch.next());
}
}
}
pub struct BatchedIter<'q, 'w, Q: Query> {
borrow: &'q mut QueryBorrow<'w, Q>,
archetype_index: u32,
batch_size: u32,
batch: u32,
}
unsafe impl<'q, 'w, Q: Query> Send for BatchedIter<'q, 'w, Q> {}
unsafe impl<'q, 'w, Q: Query> Sync for BatchedIter<'q, 'w, Q> {}
impl<'q, 'w, Q: Query> Iterator for BatchedIter<'q, 'w, Q> {
type Item = Batch<'q, Q>;
fn next(&mut self) -> Option<Self::Item> {
loop {
let archetype = self.borrow.archetypes.get(self.archetype_index as usize)?;
let offset = self.batch_size * self.batch;
if offset >= archetype.len() {
self.archetype_index += 1;
self.batch = 0;
continue;
}
if let Some(fetch) = unsafe { Q::Fetch::get(archetype, offset as usize) } {
self.batch += 1;
return Some(Batch {
_marker: PhantomData,
state: ChunkIter {
fetch,
len: self.batch_size.min(archetype.len() - offset),
},
});
} else {
self.archetype_index += 1;
debug_assert_eq!(
self.batch, 0,
"query fetch should always reject at the first batch or not at all"
);
continue;
}
}
}
}
pub struct Batch<'q, Q: Query> {
_marker: PhantomData<&'q ()>,
state: ChunkIter<Q>,
}
impl<'q, 'w, Q: Query> Iterator for Batch<'q, Q> {
type Item = <Q::Fetch as Fetch<'q>>::Item;
fn next(&mut self) -> Option<Self::Item> {
let components = unsafe { self.state.next()? };
Some(components)
}
}
unsafe impl<'q, Q: Query> Send for Batch<'q, Q> {}
unsafe impl<'q, Q: Query> Sync for Batch<'q, Q> {}
macro_rules! tuple_impl {
($($name: ident),*) => {
impl<'a, $($name: Fetch<'a>),*> Fetch<'a> for ($($name,)*) {
type Item = ($($name::Item,)*);
#[allow(unused_variables, unused_mut)]
fn access(archetype: &Archetype) -> Option<Access> {
let mut access = Access::Iterate;
$(
access = access.max($name::access(archetype)?);
)*
Some(access)
}
#[allow(unused_variables)]
fn borrow(archetype: &Archetype) {
$($name::borrow(archetype);)*
}
#[allow(unused_variables)]
unsafe fn get(archetype: &'a Archetype, offset: usize) -> Option<Self> {
Some(($($name::get(archetype, offset)?,)*))
}
#[allow(unused_variables)]
fn release(archetype: &Archetype) {
$($name::release(archetype);)*
}
#[allow(unused_variables)]
unsafe fn next(&mut self) -> Self::Item {
#[allow(non_snake_case)]
let ($($name,)*) = self;
($($name.next(),)*)
}
unsafe fn should_skip(&self) -> bool {
#[allow(non_snake_case)]
let ($($name,)*) = self;
$($name.should_skip()||)* false
}
}
impl<$($name: Query),*> Query for ($($name,)*) {
type Fetch = ($($name::Fetch,)*);
}
};
}
smaller_tuples_too!(tuple_impl, O, N, M, L, K, J, I, H, G, F, E, D, C, B, A);
#[cfg(test)]
mod tests {
use crate::{Entity, Mut, Mutated, World};
use std::{vec, vec::Vec};
use super::*;
#[test]
fn access_order() {
assert!(Access::Write > Access::Read);
assert!(Access::Read > Access::Iterate);
assert!(Some(Access::Iterate) > None);
}
struct A(usize);
struct B(usize);
struct C;
#[test]
fn added_queries() {
let mut world = World::default();
let e1 = world.spawn((A(0),));
fn get_added<Com: Component>(world: &World) -> Vec<Entity> {
world
.query::<(Added<Com>, Entity)>()
.iter()
.map(|(_added, e)| e)
.collect::<Vec<Entity>>()
};
assert_eq!(get_added::<A>(&world), vec![e1]);
world.insert(e1, (B(0),)).unwrap();
assert_eq!(get_added::<A>(&world), vec![e1]);
assert_eq!(get_added::<B>(&world), vec![e1]);
world.clear_trackers();
assert!(get_added::<A>(&world).is_empty());
let e2 = world.spawn((A(1), B(1)));
assert_eq!(get_added::<A>(&world), vec![e2]);
assert_eq!(get_added::<B>(&world), vec![e2]);
let added = world
.query::<(Entity, Added<A>, Added<B>)>()
.iter()
.map(|a| a.0)
.collect::<Vec<Entity>>();
assert_eq!(added, vec![e2]);
}
#[test]
fn mutated_trackers() {
let mut world = World::default();
let e1 = world.spawn((A(0), B(0)));
let e2 = world.spawn((A(0), B(0)));
let e3 = world.spawn((A(0), B(0)));
world.spawn((A(0), B));
for (i, mut a) in world.query::<Mut<A>>().iter().enumerate() {
if i % 2 == 0 {
a.0 += 1;
}
}
fn get_changed_a(world: &World) -> Vec<Entity> {
world
.query::<(Mutated<A>, Entity)>()
.iter()
.map(|(_a, e)| e)
.collect::<Vec<Entity>>()
};
assert_eq!(get_changed_a(&world), vec![e1, e3]);
world.insert(e1, (C,)).unwrap();
assert_eq!(get_changed_a(&world), vec![e3, e1], "changed entities list should not change (although the order will due to archetype moves)");
world.insert(e1, (A(0), B)).unwrap();
assert_eq!(
get_changed_a(&world),
vec![e3, e1],
"changed entities list should not change"
);
world.despawn(e2).unwrap();
assert_eq!(
get_changed_a(&world),
vec![e3, e1],
"changed entities list should not change"
);
world.despawn(e1).unwrap();
assert_eq!(
get_changed_a(&world),
vec![e3],
"e1 should no longer be returned"
);
world.clear_trackers();
assert!(world
.query::<(Mutated<A>, Entity)>()
.iter()
.map(|(_a, e)| e)
.collect::<Vec<Entity>>()
.is_empty());
}
#[test]
fn multiple_mutated_query() {
let mut world = World::default();
world.spawn((A(0), B(0)));
let e2 = world.spawn((A(0), B(0)));
world.spawn((A(0), B(0)));
for mut a in world.query::<Mut<A>>().iter() {
a.0 += 1;
}
for mut b in world.query::<Mut<B>>().iter().skip(1).take(1) {
b.0 += 1;
}
let a_b_changed = world
.query::<(Mutated<A>, Mutated<B>, Entity)>()
.iter()
.map(|(_a, _b, e)| e)
.collect::<Vec<Entity>>();
assert_eq!(a_b_changed, vec![e2]);
}
#[test]
fn changed_query() {
let mut world = World::default();
let e1 = world.spawn((A(0), B(0)));
fn get_changed(world: &World) -> Vec<Entity> {
world
.query::<(Changed<A>, Entity)>()
.iter()
.map(|(_a, e)| e)
.collect::<Vec<Entity>>()
};
assert_eq!(get_changed(&world), vec![e1]);
world.clear_trackers();
assert_eq!(get_changed(&world), vec![]);
*world.get_mut(e1).unwrap() = A(1);
assert_eq!(get_changed(&world), vec![e1]);
}
}