nya-ecs 0.2.0

//! Queries on [World]

use rayon::prelude::*;
use smallvec::SmallVec;
use std::{marker::PhantomData, ops::Deref};

use crate::{Component, Entity, World, component::ComponentKey};

/// A function that behaves as a predicate for a [Query]
pub trait QueryFunction: Fn(Entity, &World) -> bool + Send + Sync + 'static {}
impl<T: Fn(Entity, &World) -> bool + Send + Sync + 'static> QueryFunction for T {}

/// A filter for queries
pub struct Filter {
    include_components: Vec<ComponentKey>,
    exclude_components: Vec<ComponentKey>,
    predicate: Option<Box<dyn QueryFunction>>,
}

impl Filter {
    /// Create a [Filter]
    pub fn new() -> Self {
        Self {
            include_components: Vec::new(),
            exclude_components: Vec::new(),
            predicate: None,
        }
    }

    /// Check if `entity` satifies this filter's reqiurements
    pub fn satisfies(&self, entity: Entity, world: &World) -> bool {
        self.include_components
            .iter()
            .all(|key| world.has_key(entity, key))
            && self
                .exclude_components
                .iter()
                .all(|key| !world.has_key(entity, key))
            && (self.predicate.as_ref().is_none_or(|p| p(entity, world)))
    }

    /// Add an include set
    pub fn with_include(mut self, include: &[ComponentKey]) -> Self {
        self.include_components = include.to_vec();
        self
    }

    /// Add an exclude set
    pub fn with_exclude(mut self, exclude: &[ComponentKey]) -> Self {
        self.exclude_components = exclude.to_vec();
        self
    }

    /// Add a predicate
    pub fn with_predicate(mut self, predicate: impl QueryFunction) -> Self {
        self.predicate = Some(Box::new(predicate));
        self
    }
}

impl Default for Filter {
    fn default() -> Self {
        Self::new()
    }
}

/// Trait to convert a value to a filter
pub trait IntoFilter {
    /// Convert this value to [Filter]
    fn into_filter(self) -> Filter;
}

impl IntoFilter for Filter {
    fn into_filter(self) -> Filter {
        self
    }
}

/// [IntoFilter] is implemented for [QueryFunction]s
impl<F: QueryFunction> IntoFilter for F {
    fn into_filter(self) -> Filter {
        Filter::new().with_predicate(self)
    }
}

/// Trait to convert a type to a filter
pub trait ToFilter {
    /// Component keys of the filter
    fn keys() -> Vec<ComponentKey>;

    /// Convert to a [Filter]
    fn to_filter() -> Filter {
        Filter::new().with_include(&Self::keys())
    }
}

/// Exclude component types from a [ToFilter] tuple
#[doc(alias = "Without")]
pub struct Exclude<T: ToFilter>(PhantomData<T>);
impl<T: ToFilter> Exclude<T> {
    fn keys() -> Vec<ComponentKey> {
        T::keys()
    }
}

impl<F: ToFilter> ToFilter for Exclude<F> {
    fn keys() -> Vec<ComponentKey> {
        F::keys()
    }

    fn to_filter() -> Filter {
        Filter::new().with_exclude(&Self::keys())
    }
}

macro_rules! tuple_filter_impl {
    ($($name: ident),*) => {
        impl<$($name: Component),*> ToFilter for ($($name,)*) {
            fn keys() -> Vec<ComponentKey> {
                vec![$(ComponentKey::of::<$name>()),*]
            }
        }

        impl<$($name: Component),*, EX: ToFilter> ToFilter for ($($name,)* Exclude<EX>) {
            fn keys() -> Vec<ComponentKey> {
                vec![$(ComponentKey::of::<$name>()),*]
            }

            fn to_filter() -> Filter {
                Filter::new().with_include(&Self::keys()).with_exclude(&Exclude::<EX>::keys())
            }
        }
    };
}

nya_macros::smaller_calls_too!(
    tuple_filter_impl,
    A,
    B,
    C,
    D,
    E,
    F,
    G,
    H,
    I,
    J,
    K,
    L,
    M,
    N,
    O
);

/// A query on a [World]
#[must_use]
#[derive(Clone)]
pub struct Query {
    entities: SmallVec<[Entity; 32]>,
}

impl Query {
    /// Create a new [Query]
    pub(crate) fn new() -> Self {
        Self {
            entities: SmallVec::new(),
        }
    }

    /// Create a new [Query] from an iterator
    pub(crate) fn from_iter(iter: impl Iterator<Item = Entity>) -> Self {
        Self {
            entities: iter.collect(),
        }
    }

    /// Get the queried entities
    pub fn entities(&self) -> &[Entity] {
        &self.entities
    }

    /// Return an iterator of the query
    pub fn iter<'a>(&'a self) -> QueryIter<'a> {
        QueryIter {
            query: self,
            index: 0,
        }
    }
}

/// Query iterator
pub struct QueryIter<'a> {
    query: &'a Query,
    index: usize,
}

impl<'a> Iterator for QueryIter<'a> {
    type Item = Entity;

    fn next(&mut self) -> Option<Self::Item> {
        let index = self.index;
        self.index += 1;
        self.query.entities().get(index).copied()
    }
}

impl<'a> IntoIterator for &'a Query {
    type Item = Entity;
    type IntoIter = QueryIter<'a>;

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

impl Deref for Query {
    type Target = [Entity];

    fn deref(&self) -> &Self::Target {
        self.entities()
    }
}

impl Default for Query {
    fn default() -> Self {
        Query::new()
    }
}

mod sealed {
    pub trait Sealed {}
}
impl sealed::Sealed for crate::World {}

/// A trait for queryable objects
pub trait Queryable: sealed::Sealed {
    /// Run the query with `filters`
    fn query_filter(&self, filters: impl IntoFilter) -> Query;

    /// Run the query with filters `T`
    fn query<T: ToFilter>(&self) -> Query {
        self.query_filter(T::to_filter())
    }

    /// Query for every entity
    fn query_all(&self) -> Query;
}

impl Queryable for World {
    fn query_filter(&self, filters: impl IntoFilter) -> Query {
        let filter = filters.into_filter();
        let entities = self
            .entities
            .par_iter()
            .copied()
            .filter(|&e| filter.satisfies(e, self))
            .collect::<Vec<_>>();
        Query::from_iter(entities.into_iter())
    }

    fn query_all(&self) -> Query {
        Query::from_iter(self.entities.iter().copied())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[allow(dead_code)]
    struct Tag;
    crate::component!(Tag);
    #[allow(dead_code)]
    struct Val(u32);
    crate::component!(Val);
    #[allow(dead_code)]
    struct Point(i32, i32);
    crate::component!(Point);

    #[test]
    fn query_components() {
        let mut world = World::new();

        for i in 0..4 {
            let entity = world.spawn();
            world.add(entity, (Tag, Val(i)));
        }

        for _ in 0..6 {
            let entity = world.spawn();
            world.add(entity, (Tag, Point(3, 5)));
        }

        for e in &world.query::<(Val, Tag)>() {
            println!("Entity with `Tag` and `Val`: {e}");
        }
    }
}