nya_ecs/
query.rs

1//! Queries on [World]
2
3use rayon::prelude::*;
4use smallvec::SmallVec;
5use std::{marker::PhantomData, ops::Deref};
6
7use crate::{Component, Entity, World, component::ComponentKey};
8
9/// A function that behaves as a predicate for a [Query]
10pub trait QueryFunction: Fn(Entity, &World) -> bool + Send + Sync + 'static {}
11impl<T: Fn(Entity, &World) -> bool + Send + Sync + 'static> QueryFunction for T {}
12
13/// A filter for queries
14pub struct Filter {
15    include_components: Vec<ComponentKey>,
16    exclude_components: Vec<ComponentKey>,
17    predicate: Option<Box<dyn QueryFunction>>,
18}
19
20impl Filter {
21    /// Create a [Filter]
22    pub fn new() -> Self {
23        Self {
24            include_components: Vec::new(),
25            exclude_components: Vec::new(),
26            predicate: None,
27        }
28    }
29
30    /// Check if `entity` satifies this filter's reqiurements
31    pub fn satisfies(&self, entity: Entity, world: &World) -> bool {
32        self.include_components
33            .iter()
34            .all(|key| world.has_key(entity, key))
35            && self
36                .exclude_components
37                .iter()
38                .all(|key| !world.has_key(entity, key))
39            && (self.predicate.as_ref().is_none_or(|p| p(entity, world)))
40    }
41
42    /// Add an include set
43    pub fn with_include(mut self, include: &[ComponentKey]) -> Self {
44        self.include_components = include.to_vec();
45        self
46    }
47
48    /// Add an exclude set
49    pub fn with_exclude(mut self, exclude: &[ComponentKey]) -> Self {
50        self.exclude_components = exclude.to_vec();
51        self
52    }
53
54    /// Add a predicate
55    pub fn with_predicate(mut self, predicate: impl QueryFunction) -> Self {
56        self.predicate = Some(Box::new(predicate));
57        self
58    }
59}
60
61impl Default for Filter {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67/// Trait to convert a value to a filter
68pub trait IntoFilter {
69    /// Convert this value to [Filter]
70    fn into_filter(self) -> Filter;
71}
72
73impl IntoFilter for Filter {
74    fn into_filter(self) -> Filter {
75        self
76    }
77}
78
79/// [IntoFilter] is implemented for [QueryFunction]s
80impl<F: QueryFunction> IntoFilter for F {
81    fn into_filter(self) -> Filter {
82        Filter::new().with_predicate(self)
83    }
84}
85
86/// Trait to convert a type to a filter
87pub trait ToFilter {
88    /// Component keys of the filter
89    fn keys() -> Vec<ComponentKey>;
90
91    /// Convert to a [Filter]
92    fn to_filter() -> Filter {
93        Filter::new().with_include(&Self::keys())
94    }
95}
96
97/// Exclude component types from a [ToFilter] tuple
98#[doc(alias = "Without")]
99pub struct Exclude<T: ToFilter>(PhantomData<T>);
100impl<T: ToFilter> Exclude<T> {
101    fn keys() -> Vec<ComponentKey> {
102        T::keys()
103    }
104}
105
106impl<F: ToFilter> ToFilter for Exclude<F> {
107    fn keys() -> Vec<ComponentKey> {
108        F::keys()
109    }
110
111    fn to_filter() -> Filter {
112        Filter::new().with_exclude(&Self::keys())
113    }
114}
115
116macro_rules! tuple_filter_impl {
117    ($($name: ident),*) => {
118        impl<$($name: Component),*> ToFilter for ($($name,)*) {
119            fn keys() -> Vec<ComponentKey> {
120                vec![$(ComponentKey::of::<$name>()),*]
121            }
122        }
123
124        impl<$($name: Component),*, EX: ToFilter> ToFilter for ($($name,)* Exclude<EX>) {
125            fn keys() -> Vec<ComponentKey> {
126                vec![$(ComponentKey::of::<$name>()),*]
127            }
128
129            fn to_filter() -> Filter {
130                Filter::new().with_include(&Self::keys()).with_exclude(&Exclude::<EX>::keys())
131            }
132        }
133    };
134}
135
136nya_macros::smaller_calls_too!(
137    tuple_filter_impl,
138    A,
139    B,
140    C,
141    D,
142    E,
143    F,
144    G,
145    H,
146    I,
147    J,
148    K,
149    L,
150    M,
151    N,
152    O
153);
154
155/// A query on a [World]
156#[must_use]
157#[derive(Clone)]
158pub struct Query {
159    entities: SmallVec<[Entity; 32]>,
160}
161
162impl Query {
163    /// Create a new [Query]
164    pub(crate) fn new() -> Self {
165        Self {
166            entities: SmallVec::new(),
167        }
168    }
169
170    /// Create a new [Query] from an iterator
171    pub(crate) fn from_iter(iter: impl Iterator<Item = Entity>) -> Self {
172        Self {
173            entities: iter.collect(),
174        }
175    }
176
177    /// Get the queried entities
178    pub fn entities(&self) -> &[Entity] {
179        &self.entities
180    }
181
182    /// Return an iterator of the query
183    pub fn iter<'a>(&'a self) -> QueryIter<'a> {
184        QueryIter {
185            query: self,
186            index: 0,
187        }
188    }
189}
190
191/// Query iterator
192pub struct QueryIter<'a> {
193    query: &'a Query,
194    index: usize,
195}
196
197impl<'a> Iterator for QueryIter<'a> {
198    type Item = Entity;
199
200    fn next(&mut self) -> Option<Self::Item> {
201        let index = self.index;
202        self.index += 1;
203        self.query.entities().get(index).copied()
204    }
205}
206
207impl<'a> IntoIterator for &'a Query {
208    type Item = Entity;
209    type IntoIter = QueryIter<'a>;
210
211    fn into_iter(self) -> Self::IntoIter {
212        self.iter()
213    }
214}
215
216impl Deref for Query {
217    type Target = [Entity];
218
219    fn deref(&self) -> &Self::Target {
220        self.entities()
221    }
222}
223
224impl Default for Query {
225    fn default() -> Self {
226        Query::new()
227    }
228}
229
230mod sealed {
231    pub trait Sealed {}
232}
233impl sealed::Sealed for crate::World {}
234
235/// A trait for queryable objects
236pub trait Queryable: sealed::Sealed {
237    /// Run the query with `filters`
238    fn query_filter(&self, filters: impl IntoFilter) -> Query;
239
240    /// Run the query with filters `T`
241    fn query<T: ToFilter>(&self) -> Query {
242        self.query_filter(T::to_filter())
243    }
244
245    /// Query for every entity
246    fn query_all(&self) -> Query;
247}
248
249impl Queryable for World {
250    fn query_filter(&self, filters: impl IntoFilter) -> Query {
251        let filter = filters.into_filter();
252        let entities = self
253            .entities
254            .par_iter()
255            .copied()
256            .filter(|&e| filter.satisfies(e, self))
257            .collect::<Vec<_>>();
258        Query::from_iter(entities.into_iter())
259    }
260
261    fn query_all(&self) -> Query {
262        Query::from_iter(self.entities.iter().copied())
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[allow(dead_code)]
271    struct Tag;
272    crate::component!(Tag);
273    #[allow(dead_code)]
274    struct Val(u32);
275    crate::component!(Val);
276    #[allow(dead_code)]
277    struct Point(i32, i32);
278    crate::component!(Point);
279
280    #[test]
281    fn query_components() {
282        let mut world = World::new();
283
284        for i in 0..4 {
285            let entity = world.spawn();
286            world.add(entity, (Tag, Val(i)));
287        }
288
289        for _ in 0..6 {
290            let entity = world.spawn();
291            world.add(entity, (Tag, Point(3, 5)));
292        }
293
294        for e in &world.query::<(Val, Tag)>() {
295            println!("Entity with `Tag` and `Val`: {e}");
296        }
297    }
298}
nya_ecs/query.rs

nya_ecs/
query.rs
