use core::marker::PhantomData;
use core::ptr::NonNull;
use crate::archetype::Archetype;
use crate::entities::EntityMeta;
use crate::{Component, Entity};
pub trait Query {
#[doc(hidden)]
type Fetch: for<'a> Fetch<'a>;
}
pub type QueryItem<'a, Q> = <<Q as Query>::Fetch as Fetch<'a>>::Item;
pub unsafe trait Fetch<'a>: Sized {
type Item;
fn dangling() -> Self;
fn access(archetype: &Archetype) -> Option<Access>;
fn borrow(archetype: &Archetype);
fn new(archetype: &'a Archetype) -> Option<Self>;
fn release(archetype: &Archetype);
unsafe fn get(&self, n: usize) -> Self::Item;
}
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub enum Access {
Iterate,
Read,
Write,
}
impl<'a, T: Component> Query for &'a T {
type Fetch = FetchRead<T>;
}
#[doc(hidden)]
pub struct FetchRead<T>(NonNull<T>);
unsafe impl<'a, T: Component> Fetch<'a> for FetchRead<T> {
type Item = &'a T;
fn dangling() -> Self {
Self(NonNull::dangling())
}
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Read)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow::<T>();
}
fn new(archetype: &'a Archetype) -> Option<Self> {
archetype.get::<T>().map(Self)
}
fn release(archetype: &Archetype) {
archetype.release::<T>();
}
unsafe fn get(&self, n: usize) -> Self::Item {
&*self.0.as_ptr().add(n)
}
}
impl<'a, T: Component> Query for &'a mut T {
type Fetch = FetchWrite<T>;
}
#[doc(hidden)]
pub struct FetchWrite<T>(NonNull<T>);
unsafe impl<'a, T: Component> Fetch<'a> for FetchWrite<T> {
type Item = &'a mut T;
fn dangling() -> Self {
Self(NonNull::dangling())
}
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
Some(Access::Write)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
archetype.borrow_mut::<T>();
}
fn new(archetype: &'a Archetype) -> Option<Self> {
archetype.get::<T>().map(Self)
}
fn release(archetype: &Archetype) {
archetype.release_mut::<T>();
}
unsafe fn get(&self, n: usize) -> Self::Item {
&mut *self.0.as_ptr().add(n)
}
}
impl<T: Query> Query for Option<T> {
type Fetch = TryFetch<T::Fetch>;
}
#[doc(hidden)]
pub struct TryFetch<T>(Option<T>);
unsafe impl<'a, T: Fetch<'a>> Fetch<'a> for TryFetch<T> {
type Item = Option<T::Item>;
fn dangling() -> Self {
Self(None)
}
fn access(archetype: &Archetype) -> Option<Access> {
Some(T::access(archetype).unwrap_or(Access::Iterate))
}
fn borrow(archetype: &Archetype) {
T::borrow(archetype)
}
fn new(archetype: &'a Archetype) -> Option<Self> {
Some(Self(T::new(archetype)))
}
fn release(archetype: &Archetype) {
T::release(archetype)
}
unsafe fn get(&self, n: usize) -> Option<T::Item> {
Some(self.0.as_ref()?.get(n))
}
}
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)>);
unsafe impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWithout<T, F> {
type Item = F::Item;
fn dangling() -> Self {
Self(F::dangling(), PhantomData)
}
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
None
} else {
F::access(archetype)
}
}
fn borrow(archetype: &Archetype) {
F::borrow(archetype)
}
fn new(archetype: &'a Archetype) -> Option<Self> {
if archetype.has::<T>() {
return None;
}
Some(Self(F::new(archetype)?, PhantomData))
}
fn release(archetype: &Archetype) {
F::release(archetype)
}
unsafe fn get(&self, n: usize) -> F::Item {
self.0.get(n)
}
}
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)>);
unsafe impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWith<T, F> {
type Item = F::Item;
fn dangling() -> Self {
Self(F::dangling(), PhantomData)
}
fn access(archetype: &Archetype) -> Option<Access> {
if archetype.has::<T>() {
F::access(archetype)
} else {
None
}
}
fn borrow(archetype: &Archetype) {
F::borrow(archetype)
}
fn new(archetype: &'a Archetype) -> Option<Self> {
if !archetype.has::<T>() {
return None;
}
Some(Self(F::new(archetype)?, PhantomData))
}
fn release(archetype: &Archetype) {
F::release(archetype)
}
unsafe fn get(&self, n: usize) -> F::Item {
self.0.get(n)
}
}
pub struct QueryBorrow<'w, Q: Query> {
meta: &'w [EntityMeta],
archetypes: &'w [Archetype],
borrowed: bool,
_marker: PhantomData<Q>,
}
impl<'w, Q: Query> QueryBorrow<'w, Q> {
pub(crate) fn new(meta: &'w [EntityMeta], archetypes: &'w [Archetype]) -> Self {
Self {
meta,
archetypes,
borrowed: false,
_marker: PhantomData,
}
}
pub fn iter(&mut self) -> QueryIter<'_, Q> {
self.borrow();
unsafe { QueryIter::new(self.meta, self.archetypes) }
}
pub fn iter_batched(&mut self, batch_size: u32) -> BatchedIter<'_, Q> {
self.borrow();
unsafe { BatchedIter::new(self.meta, self.archetypes, batch_size) }
}
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 {
meta: self.meta,
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 Item = (Entity, QueryItem<'q, Q>);
type IntoIter = QueryIter<'q, Q>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
pub struct QueryIter<'q, Q: Query> {
meta: &'q [EntityMeta],
archetypes: &'q [Archetype],
archetype_index: usize,
iter: ChunkIter<Q>,
}
impl<'q, Q: Query> QueryIter<'q, Q> {
pub(crate) unsafe fn new(meta: &'q [EntityMeta], archetypes: &'q [Archetype]) -> Self {
Self {
meta,
archetypes,
archetype_index: 0,
iter: ChunkIter::empty(),
}
}
}
unsafe impl<'q, Q: Query> Send for QueryIter<'q, Q> {}
unsafe impl<'q, Q: Query> Sync for QueryIter<'q, Q> {}
impl<'q, Q: Query> Iterator for QueryIter<'q, Q> {
type Item = (Entity, QueryItem<'q, Q>);
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
loop {
match unsafe { self.iter.next() } {
None => {
let archetype = self.archetypes.get(self.archetype_index)?;
self.archetype_index += 1;
self.iter =
Q::Fetch::new(archetype).map_or(ChunkIter::empty(), |fetch| ChunkIter {
entities: archetype.entities(),
fetch,
position: 0,
len: archetype.len() as usize,
});
continue;
}
Some((id, components)) => {
return Some((
Entity {
id,
generation: unsafe { self.meta.get_unchecked(id as usize).generation },
},
components,
));
}
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let n = self.len();
(n, Some(n))
}
}
impl<'q, Q: Query> ExactSizeIterator for QueryIter<'q, Q> {
fn len(&self) -> usize {
self.archetypes
.iter()
.filter(|&x| Q::Fetch::access(x).is_some())
.map(|x| x.len() as usize)
.sum()
}
}
pub struct QueryMut<'q, Q: Query> {
iter: QueryIter<'q, Q>,
}
impl<'q, Q: Query> QueryMut<'q, Q> {
pub(crate) fn new(meta: &'q [EntityMeta], archetypes: &'q mut [Archetype]) -> Self {
Self {
iter: unsafe { QueryIter::new(meta, archetypes) },
}
}
pub fn with<T: Component>(self) -> QueryMut<'q, With<T, Q>> {
self.transform()
}
pub fn without<T: Component>(self) -> QueryMut<'q, Without<T, Q>> {
self.transform()
}
fn transform<R: Query>(self) -> QueryMut<'q, R> {
QueryMut {
iter: unsafe { QueryIter::new(self.iter.meta, self.iter.archetypes) },
}
}
}
impl<'q, Q: Query> IntoIterator for QueryMut<'q, Q> {
type Item = <QueryIter<'q, Q> as Iterator>::Item;
type IntoIter = QueryIter<'q, Q>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter
}
}
struct ChunkIter<Q: Query> {
entities: NonNull<u32>,
fetch: Q::Fetch,
position: usize,
len: usize,
}
impl<Q: Query> ChunkIter<Q> {
fn empty() -> Self {
Self {
entities: NonNull::dangling(),
fetch: Q::Fetch::dangling(),
position: 0,
len: 0,
}
}
#[inline]
unsafe fn next<'a>(&mut self) -> Option<(u32, <Q::Fetch as Fetch<'a>>::Item)> {
if self.position == self.len {
return None;
}
let entity = self.entities.as_ptr().add(self.position);
let item = self.fetch.get(self.position);
self.position += 1;
Some((*entity, item))
}
}
pub struct BatchedIter<'q, Q: Query> {
_marker: PhantomData<&'q Q>,
meta: &'q [EntityMeta],
archetypes: &'q [Archetype],
archetype_index: usize,
batch_size: u32,
batch: u32,
}
impl<'q, Q: Query> BatchedIter<'q, Q> {
pub(crate) unsafe fn new(
meta: &'q [EntityMeta],
archetypes: &'q [Archetype],
batch_size: u32,
) -> Self {
Self {
_marker: PhantomData,
meta,
archetypes,
archetype_index: 0,
batch_size,
batch: 0,
}
}
}
unsafe impl<'q, Q: Query> Send for BatchedIter<'q, Q> {}
unsafe impl<'q, Q: Query> Sync for BatchedIter<'q, Q> {}
impl<'q, Q: Query> Iterator for BatchedIter<'q, Q> {
type Item = Batch<'q, Q>;
fn next(&mut self) -> Option<Self::Item> {
loop {
let archetype = self.archetypes.get(self.archetype_index)?;
let offset = self.batch_size * self.batch;
if offset >= archetype.len() {
self.archetype_index += 1;
self.batch = 0;
continue;
}
if let Some(fetch) = Q::Fetch::new(archetype) {
self.batch += 1;
return Some(Batch {
meta: self.meta,
state: ChunkIter {
entities: archetype.entities(),
fetch,
len: (offset + self.batch_size.min(archetype.len() - offset)) as usize,
position: offset as usize,
},
});
} 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> {
meta: &'q [EntityMeta],
state: ChunkIter<Q>,
}
impl<'q, Q: Query> Iterator for Batch<'q, Q> {
type Item = (Entity, QueryItem<'q, Q>);
fn next(&mut self) -> Option<Self::Item> {
let (id, components) = unsafe { self.state.next()? };
Some((
Entity {
id,
generation: self.meta[id as usize].generation,
},
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),*) => {
unsafe impl<'a, $($name: Fetch<'a>),*> Fetch<'a> for ($($name,)*) {
type Item = ($($name::Item,)*);
fn dangling() -> Self {
($($name::dangling(),)*)
}
#[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)]
fn new(archetype: &'a Archetype) -> Option<Self> {
Some(($($name::new(archetype)?,)*))
}
#[allow(unused_variables)]
fn release(archetype: &Archetype) {
$($name::release(archetype);)*
}
#[allow(unused_variables)]
unsafe fn get(&self, n: usize) -> Self::Item {
#[allow(non_snake_case)]
let ($($name,)*) = self;
($($name.get(n),)*)
}
}
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 super::*;
#[test]
fn access_order() {
assert!(Access::Write > Access::Read);
assert!(Access::Read > Access::Iterate);
assert!(Some(Access::Iterate) > None);
}
}