bevy_defer/access/
query.rs

1use std::{
2    any::{type_name, Any, TypeId},
3    cell::RefCell,
4};
5
6use bevy::ecs::{
7    entity::{Entity, EntityEquivalent},
8    query::{
9        IterQueryData, QueryData, QueryFilter, QueryIter, QueryManyIter, QuerySingleError,
10        QueryState, ReadOnlyQueryData,
11    },
12    world::{CommandQueue, World},
13};
14use rustc_hash::FxHashMap;
15
16use crate::AccessError;
17
18/// A resource that caches a [`QueryState`].
19///
20/// Since [`QueryState`] cannot be constructed from `&World`, readonly access is not supported.
21#[derive(Debug, Default)]
22pub(crate) struct QueryCache(RefCell<FxHashMap<(TypeId, TypeId), Box<dyn Any>>>);
23
24impl QueryCache {
25    pub fn take_query_state<D: QueryData + 'static, F: QueryFilter + 'static>(
26        &self,
27    ) -> Option<Box<QueryState<D, F>>> {
28        let mut guard = self.0.borrow_mut();
29        guard
30            .remove(&(TypeId::of::<D>(), TypeId::of::<F>()))
31            .and_then(|x| x.downcast().ok())
32    }
33
34    pub fn push_query_state<D: QueryData + 'static, F: QueryFilter + 'static>(
35        &self,
36        state: Box<QueryState<D, F>>,
37    ) {
38        let mut guard = self.0.borrow_mut();
39        guard.insert((TypeId::of::<D>(), TypeId::of::<F>()), state);
40    }
41}
42
43/// A [`World`] reference with a cached [`QueryState`].
44///
45/// Tries to obtain the [`QueryState`] from the [`World`] on create
46/// and caches the [`QueryState`] in the [`World`] on drop.
47pub struct OwnedQueryState<'t, D: QueryData + 'static, F: QueryFilter + 'static> {
48    pub(crate) world: &'t mut World,
49    pub(crate) state: Option<Box<QueryState<D, F>>>,
50}
51
52impl<D: QueryData + 'static, F: QueryFilter + 'static> OwnedQueryState<'_, D, F> {
53    /// Apply a command queue to world.
54    pub fn apply_commands(&mut self, commands: &mut CommandQueue) {
55        commands.apply(self.world)
56    }
57}
58
59impl<D: QueryData + 'static, F: QueryFilter + 'static> OwnedQueryState<'_, D, F> {
60    pub fn new(world: &mut World) -> OwnedQueryState<'_, D, F> {
61        OwnedQueryState {
62            state: match world.non_send::<QueryCache>().take_query_state::<D, F>() {
63                Some(item) => Some(item),
64                None => Some(Box::new(QueryState::new(world))),
65            },
66            world,
67        }
68    }
69
70    pub fn get(
71        &mut self,
72        entity: Entity,
73    ) -> Result<<D::ReadOnly as QueryData>::Item<'_, '_>, AccessError> {
74        self.state
75            .as_mut()
76            .unwrap()
77            .get(self.world, entity)
78            .map_err(|_| AccessError::QueryConditionNotMet {
79                entity,
80                query: type_name::<(D, F)>(),
81            })
82    }
83
84    pub fn get_mut(&mut self, entity: Entity) -> Result<D::Item<'_, '_>, AccessError> {
85        self.state
86            .as_mut()
87            .unwrap()
88            .get_mut(self.world, entity)
89            .map_err(|_| AccessError::QueryConditionNotMet {
90                entity,
91                query: type_name::<(D, F)>(),
92            })
93    }
94}
95
96impl<D: IterQueryData + 'static, F: QueryFilter + 'static> OwnedQueryState<'_, D, F> {
97    pub fn single(&mut self) -> Result<<D::ReadOnly as QueryData>::Item<'_, '_>, AccessError> {
98        self.state
99            .as_mut()
100            .unwrap()
101            .single(self.world)
102            .map_err(|e| match e {
103                QuerySingleError::NoEntities(_) => AccessError::NoEntityFound {
104                    query: type_name::<D>(),
105                },
106                QuerySingleError::MultipleEntities(_) => AccessError::TooManyEntities {
107                    query: type_name::<D>(),
108                },
109            })
110    }
111
112    pub fn single_mut(&mut self) -> Result<D::Item<'_, '_>, AccessError> {
113        self.state
114            .as_mut()
115            .unwrap()
116            .single_mut(self.world)
117            .map_err(|e| match e {
118                QuerySingleError::NoEntities(_) => AccessError::NoEntityFound {
119                    query: type_name::<D>(),
120                },
121                QuerySingleError::MultipleEntities(_) => AccessError::TooManyEntities {
122                    query: type_name::<D>(),
123                },
124            })
125    }
126
127    pub fn iter_many<E: IntoIterator<Item: EntityEquivalent>>(
128        &mut self,
129        entities: E,
130    ) -> QueryManyIter<'_, '_, D::ReadOnly, F, E::IntoIter> {
131        self.state.as_mut().unwrap().iter_many(self.world, entities)
132    }
133
134    pub fn iter_many_mut<E: IntoIterator<Item: EntityEquivalent>>(
135        &mut self,
136        entities: E,
137    ) -> QueryManyIter<'_, '_, D, F, E::IntoIter> {
138        self.state
139            .as_mut()
140            .unwrap()
141            .iter_many_mut(self.world, entities)
142    }
143
144    pub fn iter(&mut self) -> QueryIter<'_, '_, D::ReadOnly, F> {
145        self.state.as_mut().unwrap().iter(self.world)
146    }
147
148    pub fn iter_mut(&mut self) -> QueryIter<'_, '_, D, F> {
149        self.state.as_mut().unwrap().iter_mut(self.world)
150    }
151}
152
153impl<'s, D: IterQueryData + 'static, F: QueryFilter + 'static> IntoIterator
154    for &'s mut OwnedQueryState<'_, D, F>
155{
156    type Item = D::Item<'s, 's>;
157    type IntoIter = QueryIter<'s, 's, D, F>;
158
159    fn into_iter(self) -> Self::IntoIter {
160        self.state.as_mut().unwrap().iter_mut(self.world)
161    }
162}
163
164impl<D: QueryData + 'static, F: QueryFilter + 'static> Drop for OwnedQueryState<'_, D, F> {
165    fn drop(&mut self) {
166        self.world
167            .non_send::<QueryCache>()
168            .push_query_state(self.state.take().unwrap());
169    }
170}
171
172/// A readonly [`World`] reference with a cached [`QueryState`].
173///
174/// # Limitation
175///
176/// Currently creates a new `QueryState` every time.
177pub struct OwnedReadonlyQueryState<'t, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> {
178    pub(crate) world: &'t World,
179    pub(crate) state: Option<Box<QueryState<D, F>>>,
180}
181
182impl<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> OwnedReadonlyQueryState<'_, D, F> {
183    pub fn new(world: &World) -> OwnedReadonlyQueryState<'_, D, F> {
184        OwnedReadonlyQueryState {
185            state: match world.non_send::<QueryCache>().take_query_state::<D, F>() {
186                Some(item) => Some(item),
187                None => QueryState::try_new(world).map(Box::new),
188            },
189            world,
190        }
191    }
192
193    pub fn single(&mut self) -> Result<D::Item<'_, '_>, AccessError> {
194        self.state
195            .as_mut()
196            .ok_or(AccessError::NoEntityFound {
197                query: type_name::<D>(),
198            })?
199            .single(self.world)
200            .map_err(|e| match e {
201                QuerySingleError::NoEntities(_) => AccessError::NoEntityFound {
202                    query: type_name::<D>(),
203                },
204                QuerySingleError::MultipleEntities(_) => AccessError::TooManyEntities {
205                    query: type_name::<D>(),
206                },
207            })
208    }
209
210    pub fn get(&mut self, entity: Entity) -> Result<D::Item<'_, '_>, AccessError> {
211        self.state
212            .as_mut()
213            .unwrap()
214            .get(self.world, entity)
215            .map_err(|_| AccessError::QueryConditionNotMet {
216                entity,
217                query: type_name::<(D, F)>(),
218            })
219    }
220
221    pub fn iter_many<E: IntoIterator<Item: EntityEquivalent>>(
222        &mut self,
223        entities: E,
224    ) -> QueryManyIter<'_, '_, D, F, E::IntoIter> {
225        self.state.as_mut().unwrap().iter_many(self.world, entities)
226    }
227
228    pub fn iter(&mut self) -> QueryIter<'_, '_, D::ReadOnly, F> {
229        self.state.as_mut().unwrap().iter(self.world)
230    }
231}
232
233impl<'s, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> IntoIterator
234    for &'s mut OwnedReadonlyQueryState<'_, D, F>
235{
236    type Item = D::Item<'s, 's>;
237    type IntoIter = QueryIter<'s, 's, D, F>;
238
239    fn into_iter(self) -> Self::IntoIter {
240        self.state.as_mut().unwrap().iter(self.world)
241    }
242}
243
244impl<D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> Drop
245    for OwnedReadonlyQueryState<'_, D, F>
246{
247    fn drop(&mut self) {
248        self.world
249            .non_send::<QueryCache>()
250            .push_query_state(self.state.take().unwrap());
251    }
252}
bevy_defer/access/query.rs

bevy_defer/access/
query.rs
