use std::any::{type_name, Any, TypeId}; use std::collections::HashSet; use std::iter::{Flatten, Map}; use std::marker::PhantomData; use std::slice::Iter as SliceIter; use std::sync::{Arc, Weak}; use crate::component::storage::Archetype; use crate::component::{ IsOptional as ComponentIsOptional, Sequence as ComponentSequence, }; use crate::lock::{Lock, ReadGuard}; use crate::system::{ NoInitParamFlag as NoInitSystemParamFlag, Param as SystemParam, System, }; use crate::{ComponentStorage, EntityComponent, WorldData}; #[derive(Debug)] pub struct Query<'world, Comps> where Comps: ComponentSequence, { component_storage: ReadGuard<'world, ComponentStorage>, component_storage_lock: Weak>, comps_pd: PhantomData, } impl<'world, Comps> Query<'world, Comps> where Comps: ComponentSequence, { #[must_use] pub fn iter<'this>(&'this self) -> ComponentIter<'world, Comps> where 'this: 'world, { #[cfg(feature = "debug")] tracing::debug!("Searching for {}", type_name::()); ComponentIter { entities: self .component_storage .find_entities(&Comps::type_ids()) .unwrap_or_else(|| panic!("Could not find {:?}", type_name::())) .iter() .map((|archetype| archetype.components.as_slice()) as ComponentIterMapFn) .flatten(), comps_pd: PhantomData, } } /// Returns a weak reference query to the same components. #[must_use] pub fn to_weak_ref(&self) -> WeakRef { WeakRef { component_storage: self.component_storage_lock.clone(), comps_pd: PhantomData, } } pub(crate) fn new(component_storage: &'world Arc>) -> Self { Self { component_storage: component_storage .read_nonblock() .expect("Failed to acquire read-only component storage lock"), component_storage_lock: Arc::downgrade(component_storage), comps_pd: PhantomData, } } } impl<'world, Comps> IntoIterator for &'world Query<'world, Comps> where Comps: ComponentSequence, { type IntoIter = ComponentIter<'world, Comps>; type Item = Comps::Refs<'world>; fn into_iter(self) -> Self::IntoIter { self.iter() } } unsafe impl<'world, Comps> SystemParam<'world> for Query<'world, Comps> where Comps: ComponentSequence, { 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_data: &'world WorldData, ) -> Self { Self::new(&world_data.component_storage) } 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_type_ids: Comps::type_ids(), }) } fn handle_pre_run(world_data: &WorldData) { let mut component_storage_lock = world_data .component_storage .write_nonblock() .expect("Failed to acquire read-write component storage lock"); #[cfg(feature = "debug")] tracing::debug!( "Adding archetypes lookup entry for components: ({})", type_name::() ); component_storage_lock.add_archetype_lookup_entry( &Comps::type_ids() .into_iter() .filter_map(|(component_id, is_optional)| { if is_optional == ComponentIsOptional::Yes { return None; } Some(component_id) }) .collect::>(), ); } } /// A entity query containing a weak reference to the world. #[derive(Debug)] pub struct WeakRef where Comps: ComponentSequence, { component_storage: Weak>, comps_pd: PhantomData, } impl WeakRef where Comps: ComponentSequence, { /// Returns a struct which can be used to retrieve a [`Query`]. /// /// Returns [`None`] if the [`World`] has been dropped. #[must_use] pub fn access(&self) -> Option> { Some(Ref { component_storage: self.component_storage.upgrade()?, _pd: PhantomData, }) } } impl Clone for WeakRef where Comps: ComponentSequence, { fn clone(&self) -> Self { Self { component_storage: self.component_storage.clone(), comps_pd: PhantomData, } } } /// Intermediate between [`Query`] and [`WeakRefQuery`]. Contains a strong reference to /// the world which is not allowed direct access to. #[derive(Debug, Clone)] pub struct Ref<'weak_ref, Comps> where Comps: ComponentSequence, { component_storage: Arc>, _pd: PhantomData<&'weak_ref Comps>, } impl<'weak_ref, Comps> Ref<'weak_ref, Comps> where Comps: ComponentSequence, { #[must_use] pub fn to_query(&self) -> Query<'_, Comps> { Query::new(&self.component_storage) } } type ComponentIterMapFn = for<'a> fn(&'a &'a Archetype) -> &'a [Vec]; pub struct ComponentIter<'world, Comps> { entities: Flatten, ComponentIterMapFn>>, comps_pd: PhantomData, } impl<'world, Comps> Iterator for ComponentIter<'world, Comps> where Comps: ComponentSequence + 'world, { type Item = Comps::Refs<'world>; fn next(&mut self) -> Option { Some(Comps::from_components(self.entities.next()?.iter())) } } #[derive(Debug)] struct QueryComponentIds { component_type_ids: Vec<(TypeId, ComponentIsOptional)>, } impl QueryComponentIds { fn contains_component_in(&self) -> bool where OtherComps: ComponentSequence, { let other_component_type_ids = OtherComps::type_ids() .into_iter() .map(|(type_id, _)| type_id) .collect::>(); self.component_type_ids .iter() .all(|(component_type_id, _)| { other_component_type_ids.contains(component_type_id) }) } }