use crate::{
archetype::{ArchetypeEntity, ArchetypeId, Archetypes},
entity::{Entities, Entity},
prelude::World,
query::{ArchetypeFilter, DebugCheckedUnwrap, QueryState, WorldQuery},
storage::{TableId, Tables},
};
use std::{borrow::Borrow, iter::FusedIterator, marker::PhantomData, mem::MaybeUninit};
use super::ReadOnlyWorldQuery;
pub struct QueryIter<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> {
tables: &'w Tables,
archetypes: &'w Archetypes,
query_state: &'s QueryState<Q, F>,
cursor: QueryIterationCursor<'w, 's, Q, F>,
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> QueryIter<'w, 's, Q, F> {
pub(crate) unsafe fn new(
world: &'w World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
QueryIter {
query_state,
tables: &world.storages().tables,
archetypes: &world.archetypes,
cursor: QueryIterationCursor::init(world, query_state, last_change_tick, change_tick),
}
}
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Iterator for QueryIter<'w, 's, Q, F> {
type Item = Q::Item<'w>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
unsafe {
self.cursor
.next(self.tables, self.archetypes, self.query_state)
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let max_size = self.cursor.max_remaining(self.tables, self.archetypes);
let archetype_query = Q::IS_ARCHETYPAL && F::IS_ARCHETYPAL;
let min_size = if archetype_query { max_size } else { 0 };
(min_size, Some(max_size))
}
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> FusedIterator for QueryIter<'w, 's, Q, F> {}
pub struct QueryManyIter<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery, I: Iterator>
where
I::Item: Borrow<Entity>,
{
entity_iter: I,
entities: &'w Entities,
tables: &'w Tables,
archetypes: &'w Archetypes,
fetch: Q::Fetch<'w>,
filter: F::Fetch<'w>,
query_state: &'s QueryState<Q, F>,
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery, I: Iterator> QueryManyIter<'w, 's, Q, F, I>
where
I::Item: Borrow<Entity>,
{
pub(crate) unsafe fn new<EntityList: IntoIterator<IntoIter = I>>(
world: &'w World,
query_state: &'s QueryState<Q, F>,
entity_list: EntityList,
last_change_tick: u32,
change_tick: u32,
) -> QueryManyIter<'w, 's, Q, F, I> {
let fetch = Q::init_fetch(
world,
&query_state.fetch_state,
last_change_tick,
change_tick,
);
let filter = F::init_fetch(
world,
&query_state.filter_state,
last_change_tick,
change_tick,
);
QueryManyIter {
query_state,
entities: &world.entities,
archetypes: &world.archetypes,
tables: &world.storages.tables,
fetch,
filter,
entity_iter: entity_list.into_iter(),
}
}
#[inline(always)]
unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<Q::Item<'w>> {
for entity in self.entity_iter.by_ref() {
let entity = *entity.borrow();
let location = match self.entities.get(entity) {
Some(location) => location,
None => continue,
};
if !self
.query_state
.matched_archetypes
.contains(location.archetype_id.index())
{
continue;
}
let archetype = self
.archetypes
.get(location.archetype_id)
.debug_checked_unwrap();
let table = self.tables.get(archetype.table_id()).debug_checked_unwrap();
Q::set_archetype(
&mut self.fetch,
&self.query_state.fetch_state,
archetype,
table,
);
F::set_archetype(
&mut self.filter,
&self.query_state.filter_state,
archetype,
table,
);
let table_row = archetype.entity_table_row(location.index);
if F::filter_fetch(&mut self.filter, entity, table_row) {
return Some(Q::fetch(&mut self.fetch, entity, table_row));
}
}
None
}
#[inline(always)]
pub fn fetch_next(&mut self) -> Option<Q::Item<'_>> {
unsafe { self.fetch_next_aliased_unchecked().map(Q::shrink) }
}
}
impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator> Iterator
for QueryManyIter<'w, 's, Q, F, I>
where
I::Item: Borrow<Entity>,
{
type Item = Q::Item<'w>;
#[inline(always)]
fn next(&mut self) -> Option<Self::Item> {
unsafe { self.fetch_next_aliased_unchecked() }
}
fn size_hint(&self) -> (usize, Option<usize>) {
let (_, max_size) = self.entity_iter.size_hint();
(0, max_size)
}
}
impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, I: Iterator> FusedIterator
for QueryManyIter<'w, 's, Q, F, I>
where
I::Item: Borrow<Entity>,
{
}
pub struct QueryCombinationIter<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery, const K: usize> {
tables: &'w Tables,
archetypes: &'w Archetypes,
query_state: &'s QueryState<Q, F>,
cursors: [QueryIterationCursor<'w, 's, Q, F>; K],
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery, const K: usize>
QueryCombinationIter<'w, 's, Q, F, K>
{
pub(crate) unsafe fn new(
world: &'w World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
let mut array: MaybeUninit<[QueryIterationCursor<'w, 's, Q, F>; K]> = MaybeUninit::uninit();
let ptr = array
.as_mut_ptr()
.cast::<QueryIterationCursor<'w, 's, Q, F>>();
if K != 0 {
ptr.write(QueryIterationCursor::init(
world,
query_state,
last_change_tick,
change_tick,
));
}
for slot in (1..K).map(|offset| ptr.add(offset)) {
slot.write(QueryIterationCursor::init_empty(
world,
query_state,
last_change_tick,
change_tick,
));
}
QueryCombinationIter {
query_state,
tables: &world.storages().tables,
archetypes: &world.archetypes,
cursors: array.assume_init(),
}
}
unsafe fn fetch_next_aliased_unchecked(&mut self) -> Option<[Q::Item<'w>; K]> {
if K == 0 {
return None;
}
'outer: for i in (0..K).rev() {
match self.cursors[i].next(self.tables, self.archetypes, self.query_state) {
Some(_) => {
for j in (i + 1)..K {
self.cursors[j] = self.cursors[j - 1].clone_cursor();
match self.cursors[j].next(self.tables, self.archetypes, self.query_state) {
Some(_) => {}
None if i > 0 => continue 'outer,
None => return None,
}
}
break;
}
None if i > 0 => continue,
None => return None,
}
}
let mut values = MaybeUninit::<[Q::Item<'w>; K]>::uninit();
let ptr = values.as_mut_ptr().cast::<Q::Item<'w>>();
for (offset, cursor) in self.cursors.iter_mut().enumerate() {
ptr.add(offset).write(cursor.peek_last().unwrap());
}
Some(values.assume_init())
}
#[inline]
pub fn fetch_next(&mut self) -> Option<[Q::Item<'_>; K]> {
unsafe {
self.fetch_next_aliased_unchecked()
.map(|array| array.map(Q::shrink))
}
}
}
impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, const K: usize> Iterator
for QueryCombinationIter<'w, 's, Q, F, K>
{
type Item = [Q::Item<'w>; K];
#[inline]
fn next(&mut self) -> Option<Self::Item> {
unsafe { QueryCombinationIter::fetch_next_aliased_unchecked(self) }
}
fn size_hint(&self) -> (usize, Option<usize>) {
fn choose(n: usize, k: usize) -> Option<usize> {
if k > n || n == 0 {
return Some(0);
}
let k = k.min(n - k);
let ks = 1..=k;
let ns = (n - k + 1..=n).rev();
ks.zip(ns)
.try_fold(1_usize, |acc, (k, n)| Some(acc.checked_mul(n)? / k))
}
let max_combinations = self
.cursors
.iter()
.enumerate()
.try_fold(0, |acc, (i, cursor)| {
let n = cursor.max_remaining(self.tables, self.archetypes);
Some(acc + choose(n, K - i)?)
});
let archetype_query = F::IS_ARCHETYPAL && Q::IS_ARCHETYPAL;
let known_max = max_combinations.unwrap_or(usize::MAX);
let min_combinations = if archetype_query { known_max } else { 0 };
(min_combinations, max_combinations)
}
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> ExactSizeIterator for QueryIter<'w, 's, Q, F>
where
F: ArchetypeFilter,
{
fn len(&self) -> usize {
self.size_hint().0
}
}
impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, const K: usize> FusedIterator
for QueryCombinationIter<'w, 's, Q, F, K>
{
}
struct QueryIterationCursor<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> {
table_id_iter: std::slice::Iter<'s, TableId>,
archetype_id_iter: std::slice::Iter<'s, ArchetypeId>,
table_entities: &'w [Entity],
archetype_entities: &'w [ArchetypeEntity],
fetch: Q::Fetch<'w>,
filter: F::Fetch<'w>,
current_len: usize,
current_index: usize,
phantom: PhantomData<Q>,
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> QueryIterationCursor<'w, 's, Q, F> {
unsafe fn clone_cursor(&self) -> Self {
Self {
table_id_iter: self.table_id_iter.clone(),
archetype_id_iter: self.archetype_id_iter.clone(),
table_entities: self.table_entities,
archetype_entities: self.archetype_entities,
fetch: Q::clone_fetch(&self.fetch),
filter: F::clone_fetch(&self.filter),
current_len: self.current_len,
current_index: self.current_index,
phantom: PhantomData,
}
}
}
impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> QueryIterationCursor<'w, 's, Q, F> {
const IS_DENSE: bool = Q::IS_DENSE && F::IS_DENSE;
unsafe fn init_empty(
world: &'w World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
QueryIterationCursor {
table_id_iter: [].iter(),
archetype_id_iter: [].iter(),
..Self::init(world, query_state, last_change_tick, change_tick)
}
}
unsafe fn init(
world: &'w World,
query_state: &'s QueryState<Q, F>,
last_change_tick: u32,
change_tick: u32,
) -> Self {
let fetch = Q::init_fetch(
world,
&query_state.fetch_state,
last_change_tick,
change_tick,
);
let filter = F::init_fetch(
world,
&query_state.filter_state,
last_change_tick,
change_tick,
);
QueryIterationCursor {
fetch,
filter,
table_entities: &[],
archetype_entities: &[],
table_id_iter: query_state.matched_table_ids.iter(),
archetype_id_iter: query_state.matched_archetype_ids.iter(),
current_len: 0,
current_index: 0,
phantom: PhantomData,
}
}
#[inline]
unsafe fn peek_last(&mut self) -> Option<Q::Item<'w>> {
if self.current_index > 0 {
let index = self.current_index - 1;
if Self::IS_DENSE {
let entity = self.table_entities.get_unchecked(index);
Some(Q::fetch(&mut self.fetch, *entity, index))
} else {
let archetype_entity = self.archetype_entities.get_unchecked(index);
Some(Q::fetch(
&mut self.fetch,
archetype_entity.entity,
archetype_entity.table_row,
))
}
} else {
None
}
}
fn max_remaining(&self, tables: &'w Tables, archetypes: &'w Archetypes) -> usize {
let remaining_matched: usize = if Self::IS_DENSE {
let ids = self.table_id_iter.clone();
ids.map(|id| tables[*id].entity_count()).sum()
} else {
let ids = self.archetype_id_iter.clone();
ids.map(|id| archetypes[*id].len()).sum()
};
remaining_matched + self.current_len - self.current_index
}
#[inline(always)]
unsafe fn next(
&mut self,
tables: &'w Tables,
archetypes: &'w Archetypes,
query_state: &'s QueryState<Q, F>,
) -> Option<Q::Item<'w>> {
if Self::IS_DENSE {
loop {
if self.current_index == self.current_len {
let table_id = self.table_id_iter.next()?;
let table = tables.get(*table_id).debug_checked_unwrap();
Q::set_table(&mut self.fetch, &query_state.fetch_state, table);
F::set_table(&mut self.filter, &query_state.filter_state, table);
self.table_entities = table.entities();
self.current_len = table.entity_count();
self.current_index = 0;
continue;
}
let entity = self.table_entities.get_unchecked(self.current_index);
if !F::filter_fetch(&mut self.filter, *entity, self.current_index) {
self.current_index += 1;
continue;
}
let item = Q::fetch(&mut self.fetch, *entity, self.current_index);
self.current_index += 1;
return Some(item);
}
} else {
loop {
if self.current_index == self.current_len {
let archetype_id = self.archetype_id_iter.next()?;
let archetype = archetypes.get(*archetype_id).debug_checked_unwrap();
let table = tables.get(archetype.table_id()).debug_checked_unwrap();
Q::set_archetype(&mut self.fetch, &query_state.fetch_state, archetype, table);
F::set_archetype(
&mut self.filter,
&query_state.filter_state,
archetype,
table,
);
self.archetype_entities = archetype.entities();
self.current_len = archetype.len();
self.current_index = 0;
continue;
}
let archetype_entity = self.archetype_entities.get_unchecked(self.current_index);
if !F::filter_fetch(
&mut self.filter,
archetype_entity.entity,
archetype_entity.table_row,
) {
self.current_index += 1;
continue;
}
let item = Q::fetch(
&mut self.fetch,
archetype_entity.entity,
archetype_entity.table_row,
);
self.current_index += 1;
return Some(item);
}
}
}
}