xuko_ecs/
query.rs

1use rayon::prelude::*;
2use std::sync::atomic::{AtomicBool, Ordering};
3
4use crate::{ComponentId, Entity, World};
5
6/// A function that checks if an entity will be filtered by the [`Query`].
7///
8/// This type is used in a [`QueryFilter`].
9pub type QueryFunction = fn(Entity, &World) -> bool;
10
11/// The [`QueryFilter`] describes the options for a [`crate::World`] query.
12#[derive(Clone, Hash)]
13pub struct QueryFilter {
14    /// Component types to include in the query
15    pub requires: Vec<ComponentId>,
16    /// Component types to exclude from the query
17    pub excludes: Vec<ComponentId>,
18    /// Whether to match only 1 entity
19    pub single: bool,
20    /// Function to do additional checks on the entity
21    pub func: Option<QueryFunction>,
22}
23
24impl QueryFilter {
25    /// Create a new empty [`QueryFilter`].
26    pub const fn new() -> Self {
27        Self {
28            requires: Vec::new(),
29            excludes: Vec::new(),
30            single: false,
31            func: None,
32        }
33    }
34
35    /// Sets [`QueryFilter::requires`] field and returns the [`QueryFilter`].
36    pub fn with_requires(mut self, requires: &[ComponentId]) -> Self {
37        self.requires = requires.to_owned();
38        self
39    }
40
41    /// Sets [`QueryFilter::excludes`] field and returns the [`QueryFilter`].
42    pub fn with_excludes(mut self, excludes: &[ComponentId]) -> Self {
43        self.excludes = excludes.to_owned();
44        self
45    }
46
47    /// Sets [`QueryFilter::single`] field and returns the [`QueryFilter`].
48    pub fn with_single(mut self, single: bool) -> Self {
49        self.single = single;
50        self
51    }
52
53    /// Sets [`QueryFilter::func`] field and returns the [`QueryFilter`].
54    pub fn with_func(mut self, func: Option<QueryFunction>) -> Self {
55        self.func = func;
56        self
57    }
58}
59
60impl Default for QueryFilter {
61    fn default() -> Self {
62        Self::new()
63    }
64}
65
66/// [`Query`] iterates over entities matching [`QueryFilter`]
67///
68/// # Examples
69///
70/// ```
71/// use xuko_ecs::prelude::*;
72///
73/// let mut world = World::new();
74///
75/// let filter = QueryFilter::new()
76///     .with_requires(&[component_id!(i32)]);
77///
78/// let queried = Query::query_world(&world, filter);
79/// // do stuff with the queried entities
80/// ```
81#[derive(Default, Clone)]
82pub struct Query {
83    queried: Vec<Entity>,
84    filter: Option<QueryFilter>,
85}
86
87impl Query {
88    /// Queries entities that match `filter` on `world`
89    #[must_use]
90    pub fn query_world(world: &World, filter: QueryFilter) -> Self {
91        let single_found = AtomicBool::new(false);
92        let queried = world
93            .entities()
94            .par_iter()
95            .copied()
96            .filter(|entity| {
97                let keys = world.get_entity(*entity).unwrap().keys();
98
99                let requires_fufilled = filter.requires.is_empty()
100                    || filter.requires.par_iter().all(|ci| keys.contains(ci));
101
102                let excludes_fulfilled = filter.excludes.is_empty()
103                    || filter.excludes.par_iter().all(|ci| !keys.contains(ci));
104
105                let func_fulfilled = if let Some(f) = filter.func.as_ref() {
106                    f(*entity, world)
107                } else {
108                    true
109                };
110
111                let fulfilled = requires_fufilled && excludes_fulfilled && func_fulfilled;
112
113                if filter.single {
114                    if single_found.load(Ordering::Acquire) {
115                        return false;
116                    }
117
118                    if fulfilled {
119                        single_found.store(true, Ordering::Release);
120                    }
121                }
122
123                fulfilled
124            })
125            .collect::<Vec<_>>();
126
127        Self {
128            queried,
129            filter: Some(filter),
130        }
131    }
132
133    /// Returns the [`QueryFilter`] used to create this [`Query`]
134    pub fn filter_used(&self) -> Option<&QueryFilter> {
135        self.filter.as_ref()
136    }
137
138    /// Return an iterator over the query
139    pub fn iter(&self) -> iter::Iter<'_> {
140        iter::Iter::new(self)
141    }
142}
143
144impl PartialEq for Query {
145    fn eq(&self, other: &Self) -> bool {
146        self.queried == other.queried
147    }
148}
149impl Eq for Query {}
150
151impl<'a> IntoIterator for &'a Query {
152    type Item = Entity;
153    type IntoIter = iter::Iter<'a>;
154
155    fn into_iter(self) -> Self::IntoIter {
156        self.iter()
157    }
158}
159
160/// Iterators for [`Query`]
161pub mod iter {
162    use super::{Entity, Query};
163
164    /// [`Query`] iterator
165    pub struct Iter<'a> {
166        iter: std::slice::Iter<'a, Entity>,
167    }
168
169    impl<'a> Iter<'a> {
170        pub(crate) fn new(query: &'a Query) -> Self {
171            Self {
172                iter: query.queried.iter(),
173            }
174        }
175    }
176
177    impl<'a> Iterator for Iter<'a> {
178        type Item = Entity;
179
180        fn next(&mut self) -> Option<Self::Item> {
181            self.iter.next().copied()
182        }
183    }
184}
xuko_ecs/query.rs

xuko_ecs/
query.rs
