path: root/ecs/src/query.rs

use std::any::type_name;
use std::borrow::Cow;
use std::marker::PhantomData;

use seq_macro::seq;

use crate::component::{Component, FromLockedOptional, Ref as ComponentRef};
use crate::entity::Handle as EntityHandle;
use crate::query::flexible::{Iter as FlexibleQueryIter, Query as FlexibleQuery};
use crate::system::{Param as SystemParam, System};
use crate::uid::Uid;
use crate::util::VecExt;
use crate::World;

pub mod flexible;
pub mod term;

#[derive(Debug)]
pub struct Query<'world, FieldTerms, FieldlessTerms = ()>
where
    FieldTerms: TermWithFieldTuple,
    FieldlessTerms: TermWithoutFieldTuple,
{
    world: &'world World,
    inner: FlexibleQuery<'world, 'static>,
    _pd: PhantomData<(FieldTerms, FieldlessTerms)>,
}

impl<'world, FieldTerms, FieldlessTerms> Query<'world, FieldTerms, FieldlessTerms>
where
    FieldTerms: TermWithFieldTuple,
    FieldlessTerms: TermWithoutFieldTuple,
{
    /// Iterates over the entities matching this query, the iterator item being the entity
    /// components.
    #[must_use]
    pub fn iter<'query>(
        &'query self,
    ) -> Iter<'query, 'world, FieldTerms, FlexibleQueryIter<'query>>
    {
        tracing::trace!("Searching for {}", std::any::type_name::<FieldTerms>());

        Iter {
            world: self.world,
            iter: self.inner.iter(),
            comps_pd: PhantomData,
        }
    }

    /// Iterates over the entities matching this query, the iterator item being the entity
    /// [`Uid`] and the matching entity components.
    #[must_use]
    pub fn iter_with_euids<'query>(
        &'query self,
    ) -> ComponentAndEuidIter<'query, 'world, FieldTerms, FlexibleQueryIter<'query>>
    {
        tracing::trace!("Searching for {}", std::any::type_name::<FieldTerms>());

        ComponentAndEuidIter {
            world: self.world,
            iter: self.inner.iter(),
            comps_pd: PhantomData,
        }
    }

    /// Iterates over the entities matching this query using the iterator returned by
    /// `func`.
    ///
    /// This function exists so that a custom [`EntityHandle`] iterator can be given to
    /// [`Iter`] without giving the user access to a reference to the [`World`].
    #[must_use]
    pub fn iter_with<'query, OutIter>(
        &'query self,
        func: impl FnOnce(FlexibleQueryIter<'query>) -> OutIter,
    ) -> Iter<'query, 'world, FieldTerms, OutIter>
    where
        OutIter: Iterator<Item = EntityHandle<'query>>,
    {
        tracing::trace!("Searching for {}", std::any::type_name::<FieldTerms>());

        Iter {
            world: self.world,
            iter: func(self.inner.iter()),
            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.inner.iter().nth(entity_index)?.uid())
    }

    pub(crate) fn new(world: &'world World) -> Self
    {
        let mut terms_builder = Terms::builder();

        FieldTerms::apply_terms_to_builder(&mut terms_builder);
        FieldlessTerms::apply_terms_to_builder(&mut terms_builder);

        Self {
            world,
            inner: world.flexible_query(terms_builder.build()),
            _pd: PhantomData,
        }
    }
}

impl<'query, 'world, FieldTerms, FieldlessTerms> IntoIterator
    for &'query Query<'world, FieldTerms, FieldlessTerms>
where
    FieldTerms: TermWithFieldTuple + 'world,
    FieldlessTerms: TermWithoutFieldTuple,
{
    type IntoIter = Iter<'query, 'world, FieldTerms, FlexibleQueryIter<'query>>;
    type Item = FieldTerms::Fields<'query>;

    fn into_iter(self) -> Self::IntoIter
    {
        self.iter()
    }
}

impl<'world, FieldTerms, FieldlessTerms> SystemParam<'world>
    for Query<'world, FieldTerms, FieldlessTerms>
where
    FieldTerms: TermWithFieldTuple,
    FieldlessTerms: TermWithoutFieldTuple,
{
    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)
    }
}

#[derive(Debug)]
pub struct Terms<'a>
{
    required_components: Cow<'a, [Uid]>,
    excluded_components: Cow<'a, [Uid]>,
}

impl<'a> Terms<'a>
{
    pub fn builder() -> TermsBuilder<'a>
    {
        TermsBuilder::default()
    }
}

#[derive(Debug, Default)]
pub struct TermsBuilder<'a>
{
    required_components: Cow<'a, [Uid]>,
    excluded_components: Cow<'a, [Uid]>,
}

pub trait TermsBuilderInterface<'a>
{
    fn with<ComponentT: Component>(self) -> Self;

    fn without<ComponentT: Component>(self) -> Self;

    fn with_required_ids(self, ids: &'a mut [Uid]) -> Self;
}

macro_rules! impl_terms_builder {
    ($($impl_content: tt)*) => {
        impl<'a> TermsBuilderInterface<'a> for TermsBuilder<'a> {
            $($impl_content)*
        }

        impl<'a> TermsBuilderInterface<'a> for &mut TermsBuilder<'a> {
            $($impl_content)*
        }
    };
}

impl_terms_builder! {
    #[allow(unused_mut)]
    fn with<ComponentT: Component>(mut self) -> Self
    {
        if ComponentT::is_optional() {
            return self;
        }

        self.required_components
            .to_mut()
            .insert_at_partition_point_by_key(ComponentT::id(), |id| *id);

        self
    }

    #[allow(unused_mut)]
    fn without<ComponentT: Component>(mut self) -> Self
    {
        if ComponentT::is_optional() {
            panic!(
                "{}::without cannot take optional component",
                type_name::<Self>()
            );
        }

        self.excluded_components
            .to_mut()
            .insert_at_partition_point_by_key(ComponentT::id(), |id| *id);

        self
    }

    #[allow(unused_mut)]
    fn with_required_ids(mut self, ids: &'a mut [Uid]) -> Self
    {
        if ids.is_empty() {
            return self;
        }

        if !ids.is_sorted() {
            ids.sort();
        }

        if self.required_components.is_empty() {
            self.required_components = Cow::Borrowed(ids);

            return self;
        }

        let first_id_pp_index = self.required_components.partition_point(|req_comp_id| {
            req_comp_id <= ids.first().expect("Cannot happend since not empty")
        });

        let removed = self
            .required_components
            .to_mut()
            .splice(first_id_pp_index..first_id_pp_index, ids.iter().copied());

        assert_eq!(removed.count(), 0);

        self
    }
}

impl<'a> TermsBuilder<'a>
{
    pub fn build(self) -> Terms<'a>
    {
        assert!(self.required_components.is_sorted());
        assert!(self.excluded_components.is_sorted());

        Terms {
            required_components: self.required_components,
            excluded_components: self.excluded_components,
        }
    }
}

pub trait TermWithoutField
{
    fn apply_to_terms_builder(terms_builder: &mut TermsBuilder<'_>);
}

pub trait TermWithField
{
    type Field<'a>;

    fn apply_to_terms_builder(terms_builder: &mut TermsBuilder<'_>);

    fn get_field<'world>(
        entity_handle: &EntityHandle<'world>,
        world: &'world World,
    ) -> Self::Field<'world>;
}

impl<ComponentRefT: ComponentRef> TermWithField for ComponentRefT
{
    type Field<'a> = ComponentRefT::Handle<'a>;

    fn apply_to_terms_builder(terms_builder: &mut TermsBuilder<'_>)
    {
        terms_builder.with::<ComponentRefT::Component>();
    }

    fn get_field<'world>(
        entity_handle: &EntityHandle<'world>,
        world: &'world World,
    ) -> Self::Field<'world>
    {
        Self::Field::from_locked_optional_component(
            entity_handle
                .get_component(ComponentRefT::Component::id())
                .map(|component| &component.component),
            world,
        )
        .unwrap_or_else(|err| {
            panic!(
                "Taking component {} lock failed: {err}",
                type_name::<ComponentRefT::Component>()
            );
        })
    }
}

pub trait TermWithoutFieldTuple
{
    fn apply_terms_to_builder(terms_builder: &mut TermsBuilder<'_>);
}

pub trait TermWithFieldTuple
{
    type Fields<'component>;

    fn apply_terms_to_builder(terms_builder: &mut TermsBuilder<'_>);

    fn get_fields<'component>(
        entity_handle: &EntityHandle<'component>,
        world: &'component World,
    ) -> Self::Fields<'component>;
}

pub struct Iter<'query, 'world, FieldTerms, EntityHandleIter>
where
    FieldTerms: TermWithFieldTuple + 'world,
    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
{
    world: &'world World,
    iter: EntityHandleIter,
    comps_pd: PhantomData<FieldTerms>,
}

impl<'query, 'world, FieldTerms, EntityHandleIter>
    Iter<'query, 'world, FieldTerms, EntityHandleIter>
where
    FieldTerms: TermWithFieldTuple + 'world,
    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
    'world: 'query,
{
    /// Creates a new iterator from the given entity handle iterator.
    ///
    /// # Important
    /// All of the yielded entities of the entity handle iterator should match the
    /// terms `Terms`. The [`Self::next`] function will panic if it encounters a
    /// entity that does not match the terms `Terms`.
    pub fn new(world: &'world World, iter: EntityHandleIter) -> Self
    {
        Self { world, iter, comps_pd: PhantomData }
    }
}

impl<'query, 'world, FieldTerms, EntityHandleIter> Iterator
    for Iter<'query, 'world, FieldTerms, EntityHandleIter>
where
    FieldTerms: TermWithFieldTuple + 'world,
    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
    'world: 'query,
{
    type Item = FieldTerms::Fields<'query>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let entity_handle = self.iter.next()?;

        Some(FieldTerms::get_fields(&entity_handle, self.world))
    }
}

pub struct ComponentAndEuidIter<'query, 'world, FieldTerms, EntityHandleIter>
where
    FieldTerms: TermWithFieldTuple + 'world,
    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
{
    world: &'world World,
    iter: EntityHandleIter,
    comps_pd: PhantomData<FieldTerms>,
}

impl<'query, 'world, FieldTerms, EntityHandleIter> Iterator
    for ComponentAndEuidIter<'query, 'world, FieldTerms, EntityHandleIter>
where
    FieldTerms: TermWithFieldTuple + 'world,
    EntityHandleIter: Iterator<Item = EntityHandle<'query>>,
    'world: 'query,
{
    type Item = (Uid, FieldTerms::Fields<'query>);

    fn next(&mut self) -> Option<Self::Item>
    {
        let entity_handle = self.iter.next()?;

        Some((
            entity_handle.uid(),
            FieldTerms::get_fields(&entity_handle, self.world),
        ))
    }
}

macro_rules! impl_term_sequence {
    ($c: tt) => {
        seq!(I in 0..=$c {
            impl<#(Term~I: TermWithoutField,)*> TermWithoutFieldTuple for (#(Term~I,)*)
            {
                fn apply_terms_to_builder(terms_builder: &mut TermsBuilder<'_>)
                {
                    #(
                        Term~I::apply_to_terms_builder(terms_builder);
                    )*
                }
            }

            impl<#(Term~I: TermWithField,)*> TermWithFieldTuple for (#(Term~I,)*)
            {
                type Fields<'component> = (#(Term~I::Field<'component>,)*);

                fn apply_terms_to_builder(terms_builder: &mut TermsBuilder<'_>)
                {
                    #(
                        Term~I::apply_to_terms_builder(terms_builder);
                    )*
                }

                fn get_fields<'component>(
                    entity_handle: &EntityHandle<'component>,
                    world: &'component World,
                ) -> Self::Fields<'component>
                {
                    (#(Term~I::get_field(entity_handle, world),)*)
                }
            }
        });
    };
}

seq!(C in 0..=16 {
    impl_term_sequence!(C);
});

impl TermWithoutFieldTuple for ()
{
    fn apply_terms_to_builder(_terms_builder: &mut TermsBuilder<'_>) {}
}

impl TermWithFieldTuple for ()
{
    type Fields<'component> = ();

    fn apply_terms_to_builder(_terms_builder: &mut TermsBuilder<'_>) {}

    fn get_fields<'component>(
        _entity_handle: &EntityHandle<'_>,
        _world: &'component World,
    ) -> Self::Fields<'component>
    {
    }
}