path: root/ecs/src/query.rs

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::lock::{ReadGuard, WriteGuard};
use crate::query::options::Options;
use crate::system::{
    NoInitParamFlag as NoInitSystemParamFlag,
    Param as SystemParam,
    System,
};
use crate::uid::Uid;
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::<Comps>());

        #[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::<Comps>());

        #[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,
        }
    }

    /// Iterates over the entities matching this query and has the provided extra
    /// component.
    #[must_use]
    pub fn iter_with_extra_comps(
        &'world self,
        extra_components: impl IntoIterator<Item = ComponentMetadata>,
    ) -> ComponentIter<'world, Comps, QueryEntityIter<'world>>
    {
        #[cfg(feature = "debug")]
        tracing::debug!(
            "Searching for {} + extra components",
            std::any::type_name::<Comps>()
        );

        #[allow(clippy::map_flatten)]
        ComponentIter {
            world: self.world,
            entities: self
                .component_storage
                .find_entities(
                    Comps::metadata()
                        .into_iter()
                        .chain(extra_components)
                        .collect::<Vec<_>>(),
                )
                .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 get_entity_uid(&self, entity_index: usize) -> Option<Uid>
    {
        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()
    }
}

impl<'world, Comps, OptionsT> SystemParam<'world> for Query<'world, Comps, OptionsT>
where
    Comps: ComponentSequence,
    OptionsT: Options,
{
    type Flags = NoInitSystemParamFlag;
    type Input = ();

    fn initialize<SystemImpl>(
        _system: &mut impl System<'world, SystemImpl>,
        _input: Self::Input,
    )
    {
    }

    fn new<SystemImpl>(
        _system: &'world impl System<'world, SystemImpl>,
        world: &'world World,
    ) -> Self
    {
        Self::new(world)
    }
}

type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> EntityIter<'a>;

type ComponentIterFilterFn = for<'a, 'b> fn(&'a &'b ArchetypeEntity) -> bool;

type QueryEntityIter<'world> = Filter<
    Flatten<Map<ArchetypeRefIter<'world>, ComponentIterMapFn>>,
    ComponentIterFilterFn,
>;

pub struct ComponentIterMut<'world, Comps, EntityIter>
where
    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
{
    world: &'world World,
    entities: EntityIter,
    comps_pd: PhantomData<Comps>,
}

impl<'world, Comps, EntityIter> Iterator for ComponentIterMut<'world, Comps, EntityIter>
where
    Comps: ComponentSequence + 'world,
    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
{
    type Item = Comps::MutRefs<'world>;

    fn next(&mut self) -> Option<Self::Item>
    {
        Some(Comps::from_components_mut(
            self.entities.next()?.components().iter(),
            self.world,
            lock_component_rw,
        ))
    }
}

fn lock_component_rw(
    entity_component: &EntityComponent,
) -> WriteGuard<'_, Box<dyn Component>>
{
    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<Item = &'world ArchetypeEntity>,
{
    world: &'world World,
    entities: EntityIter,
    comps_pd: PhantomData<Comps>,
}

impl<'world, Comps, EntityIter> Iterator for ComponentIter<'world, Comps, EntityIter>
where
    Comps: ComponentSequence + 'world,
    EntityIter: Iterator<Item = &'world ArchetypeEntity>,
{
    type Item = Comps::Refs<'world>;

    fn next(&mut self) -> Option<Self::Item>
    {
        Some(Comps::from_components(
            self.entities.next()?.components().iter(),
            self.world,
            lock_component_ro,
        ))
    }
}

fn lock_component_ro(
    entity_component: &EntityComponent,
) -> ReadGuard<'_, Box<dyn Component>>
{
    entity_component
        .component
        .read_nonblock()
        .unwrap_or_else(|_| {
            panic!(
                "Failed to acquire read-write lock to component {}",
                entity_component.name
            );
        })
}