use std::any::Any; use std::collections::HashSet; use std::iter::{Filter, Flatten, Map}; use std::marker::PhantomData; use crate::component::storage::{ Archetype, ArchetypeEntity, ArchetypeRefIter, EntityIter, Storage as ComponentStorage, }; use crate::component::{ Component, Metadata as ComponentMetadata, Sequence as ComponentSequence, }; use crate::entity::Uid as EntityUid; use crate::lock::{ReadGuard, WriteGuard}; use crate::query::options::Options; use crate::system::{ NoInitParamFlag as NoInitSystemParamFlag, Param as SystemParam, System, }; use crate::{EntityComponent, World}; pub mod options; #[derive(Debug)] pub struct Query<'world, Comps, OptionsT = ()> where Comps: ComponentSequence, { world: &'world World, component_storage: ReadGuard<'world, ComponentStorage>, _pd: PhantomData<(Comps, OptionsT)>, } impl<'world, Comps, OptionsT> Query<'world, Comps, OptionsT> where Comps: ComponentSequence, OptionsT: Options, { /// Iterates over the entities matching this query. #[must_use] pub fn iter_mut( &'world self, ) -> ComponentIterMut<'world, Comps, QueryEntityIter<'world>> { #[cfg(feature = "debug")] tracing::debug!("Searching for {}", std::any::type_name::()); #[allow(clippy::map_flatten)] ComponentIterMut { world: self.world, entities: self .component_storage .find_entities(Comps::metadata()) .map(Archetype::entities as ComponentIterMapFn) .flatten() .filter(|entity| OptionsT::entity_filter(entity.components())), comps_pd: PhantomData, } } /// Iterates over the entities matching this query. #[must_use] pub fn iter(&'world self) -> ComponentIter<'world, Comps, QueryEntityIter<'world>> { #[cfg(feature = "debug")] tracing::debug!("Searching for {}", std::any::type_name::()); #[allow(clippy::map_flatten)] ComponentIter { world: self.world, entities: self .component_storage .find_entities(Comps::metadata()) .map(Archetype::entities as ComponentIterMapFn) .flatten() .filter(|entity| OptionsT::entity_filter(entity.components())), comps_pd: PhantomData, } } /// Returns the UID of the entity at the given query iteration index. #[must_use] pub fn entity_uid(&self, entity_index: usize) -> Option { Some( self.component_storage .find_entities(Comps::metadata()) .flat_map(|archetype| archetype.entities()) .filter(|entity| OptionsT::entity_filter(entity.components())) .nth(entity_index)? .uid(), ) } pub(crate) fn new(world: &'world World) -> Self { Self { world, component_storage: world .data .component_storage .read_nonblock() .expect("Failed to acquire read-only component storage lock"), _pd: PhantomData, } } } impl<'world, Comps, OptionsT> IntoIterator for &'world Query<'world, Comps, OptionsT> where Comps: ComponentSequence, OptionsT: Options, { type IntoIter = ComponentIterMut<'world, Comps, QueryEntityIter<'world>>; type Item = Comps::MutRefs<'world>; fn into_iter(self) -> Self::IntoIter { self.iter_mut() } } unsafe impl<'world, Comps, OptionsT> SystemParam<'world> for Query<'world, Comps, OptionsT> where Comps: ComponentSequence, OptionsT: Options, { type Flags = NoInitSystemParamFlag; type Input = (); fn initialize( _system: &mut impl System<'world, SystemImpl>, _input: Self::Input, ) { } fn new( _system: &'world impl System<'world, SystemImpl>, world: &'world World, ) -> Self { Self::new(world) } fn is_compatible>() -> bool { let other_comparable = Other::get_comparable(); let Some(other_query_component_ids) = other_comparable.downcast_ref::() else { return true; }; !other_query_component_ids.contains_component_in::() } fn get_comparable() -> Box { Box::new(QueryComponentIds { component_ids: Comps::metadata() }) } } type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> EntityIter<'a>; type ComponentIterFilterFn = for<'a, 'b> fn(&'a &'b ArchetypeEntity) -> bool; type QueryEntityIter<'world> = Filter< Flatten, ComponentIterMapFn>>, ComponentIterFilterFn, >; pub struct ComponentIterMut<'world, Comps, EntityIter> where EntityIter: Iterator, { world: &'world World, entities: EntityIter, comps_pd: PhantomData, } impl<'world, Comps, EntityIter> Iterator for ComponentIterMut<'world, Comps, EntityIter> where Comps: ComponentSequence + 'world, EntityIter: Iterator, { type Item = Comps::MutRefs<'world>; fn next(&mut self) -> Option { Some(Comps::from_components_mut( self.entities.next()?.components().iter(), self.world, lock_component_rw, )) } } fn lock_component_rw( entity_component: &EntityComponent, ) -> WriteGuard<'_, Box> { entity_component .component .write_nonblock() .unwrap_or_else(|_| { panic!( "Failed to acquire read-write lock to component {}", entity_component.name ); }) } pub struct ComponentIter<'world, Comps, EntityIter> where EntityIter: Iterator, { world: &'world World, entities: EntityIter, comps_pd: PhantomData, } impl<'world, Comps, EntityIter> Iterator for ComponentIter<'world, Comps, EntityIter> where Comps: ComponentSequence + 'world, EntityIter: Iterator, { type Item = Comps::Refs<'world>; fn next(&mut self) -> Option { Some(Comps::from_components( self.entities.next()?.components().iter(), self.world, lock_component_ro, )) } } fn lock_component_ro( entity_component: &EntityComponent, ) -> ReadGuard<'_, Box> { entity_component .component .read_nonblock() .unwrap_or_else(|_| { panic!( "Failed to acquire read-write lock to component {}", entity_component.name ); }) } #[derive(Debug)] struct QueryComponentIds { component_ids: Vec, } impl QueryComponentIds { fn contains_component_in(&self) -> bool where OtherComps: ComponentSequence, { let other_ids = OtherComps::metadata() .into_iter() .map(|component_metadata| component_metadata.id) .collect::>(); self.component_ids .iter() .all(|component_metadata| other_ids.contains(&component_metadata.id)) } }