Skip to content

Commit fba0545

Browse files
committed
query filters
1 parent 60fa2d5 commit fba0545

File tree

21 files changed

+470
-558
lines changed

21 files changed

+470
-558
lines changed

crates/bevy_ecs/hecs/macros/src/lib.rs

+9-6
Original file line numberDiff line numberDiff line change
@@ -281,38 +281,41 @@ pub fn impl_query_set(_input: TokenStream) -> TokenStream {
281281
let mut tokens = TokenStream::new();
282282
let max_queries = 4;
283283
let queries = get_idents(|i| format!("Q{}", i), max_queries);
284+
let filters = get_idents(|i| format!("F{}", i), max_queries);
284285
let lifetimes = get_lifetimes(|i| format!("'q{}", i), max_queries);
285286
let mut query_fns = Vec::new();
286287
let mut query_fn_muts = Vec::new();
287288
for i in 0..max_queries {
288289
let query = &queries[i];
289290
let lifetime = &lifetimes[i];
291+
let filter = &filters[i];
290292
let fn_name = Ident::new(&format!("q{}", i), Span::call_site());
291293
let fn_name_mut = Ident::new(&format!("q{}_mut", i), Span::call_site());
292294
let index = Index::from(i);
293295
query_fns.push(quote! {
294-
pub fn #fn_name(&self) -> &Query<#lifetime, #query> {
296+
pub fn #fn_name(&self) -> &Query<#lifetime, #query, #filter> {
295297
&self.value.#index
296298
}
297299
});
298300
query_fn_muts.push(quote! {
299-
pub fn #fn_name_mut(&mut self) -> &mut Query<#lifetime, #query> {
301+
pub fn #fn_name_mut(&mut self) -> &mut Query<#lifetime, #query, #filter> {
300302
&mut self.value.#index
301303
}
302304
});
303305
}
304306

305307
for query_count in 1..=max_queries {
306308
let query = &queries[0..query_count];
309+
let filter = &filters[0..query_count];
307310
let lifetime = &lifetimes[0..query_count];
308311
let query_fn = &query_fns[0..query_count];
309312
let query_fn_mut = &query_fn_muts[0..query_count];
310313
tokens.extend(TokenStream::from(quote! {
311-
impl<#(#lifetime,)* #(#query: HecsQuery,)*> QueryTuple for (#(Query<#lifetime, #query>,)*) {
314+
impl<#(#lifetime,)* #(#query: HecsQuery,)* #(#filter: QueryFilter,)*> QueryTuple for (#(Query<#lifetime, #query, #filter>,)*) {
312315
unsafe fn new(world: &World, component_access: &TypeAccess<ArchetypeComponent>) -> Self {
313316
(
314317
#(
315-
Query::<#query>::new(
318+
Query::<#query, #filter>::new(
316319
std::mem::transmute(world),
317320
std::mem::transmute(component_access),
318321
),
@@ -322,12 +325,12 @@ pub fn impl_query_set(_input: TokenStream) -> TokenStream {
322325

323326
fn get_accesses() -> Vec<QueryAccess> {
324327
vec![
325-
#(<#query::Fetch as Fetch>::access(),)*
328+
#(QueryAccess::union(vec![<#query::Fetch as Fetch>::access(), #filter::access()]),)*
326329
]
327330
}
328331
}
329332

330-
impl<#(#lifetime,)* #(#query: HecsQuery,)*> QuerySet<(#(Query<#lifetime, #query>,)*)> {
333+
impl<#(#lifetime,)* #(#query: HecsQuery,)* #(#filter: QueryFilter,)*> QuerySet<(#(Query<#lifetime, #query, #filter>,)*)> {
331334
#(#query_fn)*
332335
#(#query_fn_mut)*
333336
}

crates/bevy_ecs/hecs/src/access.rs

+4-3
Original file line numberDiff line numberDiff line change
@@ -282,10 +282,11 @@ impl<T: Hash + Eq + PartialEq + Copy> TypeAccess<T> {
282282

283283
#[cfg(test)]
284284
mod tests {
285-
use crate::{ArchetypeComponent, Entity, Fetch, Query, TypeAccess, With, World};
285+
use crate::{ArchetypeComponent, Entity, Fetch, Query, QueryAccess, TypeAccess, World};
286286
use std::vec;
287287

288288
struct A;
289+
#[derive(Clone, Eq, PartialEq, Debug)]
289290
struct B;
290291
struct C;
291292

@@ -344,7 +345,7 @@ mod tests {
344345
);
345346

346347
let mut a_with_b_type_access = TypeAccess::default();
347-
<With<B, &A> as Query>::Fetch::access()
348+
QueryAccess::with::<B>(<&A as Query>::Fetch::access())
348349
.get_world_archetype_access(&world, Some(&mut a_with_b_type_access));
349350

350351
assert_eq!(
@@ -353,7 +354,7 @@ mod tests {
353354
);
354355

355356
let mut a_with_b_option_c_type_access = TypeAccess::default();
356-
<With<B, (&A, Option<&mut C>)> as Query>::Fetch::access()
357+
QueryAccess::with::<B>(<(&A, Option<&mut C>) as Query>::Fetch::access())
357358
.get_world_archetype_access(&world, Some(&mut a_with_b_option_c_type_access));
358359

359360
assert_eq!(

crates/bevy_ecs/hecs/src/filter.rs

+245
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,245 @@
1+
use crate::{archetype::Archetype, Component, QueryAccess};
2+
use core::{any::TypeId, marker::PhantomData, ptr::NonNull};
3+
use std::{boxed::Box, vec};
4+
5+
pub trait QueryFilter: Sized {
6+
type EntityFilter: EntityFilter;
7+
fn access() -> QueryAccess;
8+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter>;
9+
}
10+
11+
pub trait EntityFilter: Sized {
12+
const DANGLING: Self;
13+
14+
/// # Safety
15+
/// This might access archetype data in an unsafe manner. In general filters should be read-only and they should only access
16+
/// the data they have claimed in `access()`.
17+
unsafe fn matches_entity(&self, _offset: usize) -> bool;
18+
}
19+
20+
pub struct AnyEntityFilter;
21+
22+
impl EntityFilter for AnyEntityFilter {
23+
const DANGLING: Self = AnyEntityFilter;
24+
25+
#[inline]
26+
unsafe fn matches_entity(&self, _offset: usize) -> bool {
27+
true
28+
}
29+
}
30+
31+
pub struct Or<T>(pub T);
32+
33+
/// Query transformer that retrieves components of type `T` that have been mutated since the start of the frame.
34+
/// Added components do not count as mutated.
35+
pub struct Mutated<T>(NonNull<bool>, PhantomData<T>);
36+
37+
/// Query transformer that retrieves components of type `T` that have been added since the start of the frame.
38+
pub struct Added<T>(NonNull<bool>, PhantomData<T>);
39+
40+
/// Query transformer that retrieves components of type `T` that have either been mutated or added since the start of the frame.
41+
pub struct Changed<T>(NonNull<bool>, NonNull<bool>, PhantomData<T>);
42+
43+
impl QueryFilter for () {
44+
type EntityFilter = AnyEntityFilter;
45+
46+
fn access() -> QueryAccess {
47+
QueryAccess::None
48+
}
49+
50+
#[inline]
51+
fn get_entity_filter(_archetype: &Archetype) -> Option<Self::EntityFilter> {
52+
Some(AnyEntityFilter)
53+
}
54+
}
55+
56+
impl<T: Component> QueryFilter for Added<T> {
57+
type EntityFilter = Self;
58+
59+
fn access() -> QueryAccess {
60+
QueryAccess::read::<T>()
61+
}
62+
63+
#[inline]
64+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
65+
archetype
66+
.get_type_state(TypeId::of::<T>())
67+
.map(|state| Added(state.added(), Default::default()))
68+
}
69+
}
70+
71+
impl<T: Component> EntityFilter for Added<T> {
72+
const DANGLING: Self = Added(NonNull::dangling(), PhantomData::<T>);
73+
74+
#[inline]
75+
unsafe fn matches_entity(&self, offset: usize) -> bool {
76+
*self.0.as_ptr().add(offset)
77+
}
78+
}
79+
80+
impl<T: Component> QueryFilter for Mutated<T> {
81+
type EntityFilter = Self;
82+
83+
fn access() -> QueryAccess {
84+
QueryAccess::read::<T>()
85+
}
86+
87+
#[inline]
88+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
89+
archetype
90+
.get_type_state(TypeId::of::<T>())
91+
.map(|state| Mutated(state.mutated(), Default::default()))
92+
}
93+
}
94+
95+
impl<T: Component> EntityFilter for Mutated<T> {
96+
const DANGLING: Self = Mutated(NonNull::dangling(), PhantomData::<T>);
97+
98+
unsafe fn matches_entity(&self, offset: usize) -> bool {
99+
*self.0.as_ptr().add(offset)
100+
}
101+
}
102+
103+
impl<T: Component> QueryFilter for Changed<T> {
104+
type EntityFilter = Self;
105+
106+
fn access() -> QueryAccess {
107+
QueryAccess::read::<T>()
108+
}
109+
110+
#[inline]
111+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
112+
archetype
113+
.get_type_state(TypeId::of::<T>())
114+
.map(|state| Changed(state.added(), state.mutated(), Default::default()))
115+
}
116+
}
117+
118+
impl<T: Component> EntityFilter for Changed<T> {
119+
const DANGLING: Self = Changed(NonNull::dangling(), NonNull::dangling(), PhantomData::<T>);
120+
121+
#[inline]
122+
unsafe fn matches_entity(&self, offset: usize) -> bool {
123+
*self.0.as_ptr().add(offset) || *self.1.as_ptr().add(offset)
124+
}
125+
}
126+
127+
pub struct Without<T>(PhantomData<T>);
128+
129+
impl<T: Component> QueryFilter for Without<T> {
130+
type EntityFilter = AnyEntityFilter;
131+
132+
fn access() -> QueryAccess {
133+
QueryAccess::without::<T>(QueryAccess::None)
134+
}
135+
136+
#[inline]
137+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
138+
if archetype.has_type(TypeId::of::<T>()) {
139+
None
140+
} else {
141+
Some(AnyEntityFilter)
142+
}
143+
}
144+
}
145+
146+
pub struct With<T>(PhantomData<T>);
147+
148+
impl<T: Component> QueryFilter for With<T> {
149+
type EntityFilter = AnyEntityFilter;
150+
151+
fn access() -> QueryAccess {
152+
QueryAccess::with::<T>(QueryAccess::None)
153+
}
154+
155+
#[inline]
156+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
157+
if archetype.has_type(TypeId::of::<T>()) {
158+
Some(AnyEntityFilter)
159+
} else {
160+
None
161+
}
162+
}
163+
}
164+
165+
macro_rules! impl_query_filter_tuple {
166+
($($filter: ident),*) => {
167+
#[allow(unused_variables)]
168+
#[allow(non_snake_case)]
169+
impl<$($filter: QueryFilter),*> QueryFilter for ($($filter,)*) {
170+
type EntityFilter = ($($filter::EntityFilter,)*);
171+
172+
fn access() -> QueryAccess {
173+
QueryAccess::union(vec![
174+
$($filter::access(),)+
175+
])
176+
}
177+
178+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
179+
Some(($($filter::get_entity_filter(archetype)?,)*))
180+
}
181+
182+
}
183+
184+
#[allow(unused_variables)]
185+
#[allow(non_snake_case)]
186+
impl<$($filter: EntityFilter),*> EntityFilter for ($($filter,)*) {
187+
const DANGLING: Self = ($($filter::DANGLING,)*);
188+
unsafe fn matches_entity(&self, offset: usize) -> bool {
189+
let ($($filter,)*) = self;
190+
true $(&& $filter.matches_entity(offset))*
191+
}
192+
}
193+
194+
#[allow(unused_variables)]
195+
#[allow(non_snake_case)]
196+
impl<$($filter: QueryFilter),*> QueryFilter for Or<($($filter,)*)> {
197+
type EntityFilter = Or<($(Option<<$filter as QueryFilter>::EntityFilter>,)*)>;
198+
fn access() -> QueryAccess {
199+
QueryAccess::union(vec![
200+
$(QueryAccess::Optional(Box::new($filter::access())),)+
201+
])
202+
}
203+
204+
fn get_entity_filter(archetype: &Archetype) -> Option<Self::EntityFilter> {
205+
let mut matches_something = false;
206+
$(
207+
let $filter = $filter::get_entity_filter(archetype);
208+
matches_something = matches_something || $filter.is_some();
209+
)*
210+
if matches_something {
211+
Some(Or(($($filter,)*)))
212+
} else {
213+
None
214+
}
215+
}
216+
217+
}
218+
#[allow(unused_variables)]
219+
#[allow(non_snake_case)]
220+
impl<$($filter: EntityFilter),*> EntityFilter for Or<($(Option<$filter>,)*)> {
221+
const DANGLING: Self = Or(($(Some($filter::DANGLING),)*));
222+
unsafe fn matches_entity(&self, offset: usize) -> bool {
223+
let Or(($($filter,)*)) = self;
224+
false $(|| $filter.as_ref().map_or(false, |filter|filter.matches_entity(offset)))*
225+
}
226+
}
227+
};
228+
}
229+
230+
impl_query_filter_tuple!(A);
231+
impl_query_filter_tuple!(A, B);
232+
impl_query_filter_tuple!(A, B, C);
233+
impl_query_filter_tuple!(A, B, C, D);
234+
impl_query_filter_tuple!(A, B, C, D, E);
235+
impl_query_filter_tuple!(A, B, C, D, E, F);
236+
impl_query_filter_tuple!(A, B, C, D, E, F, G);
237+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H);
238+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I);
239+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J);
240+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K);
241+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L);
242+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M);
243+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, N);
244+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O);
245+
impl_query_filter_tuple!(A, B, C, D, E, F, G, H, I, J, K, L, M, N, O, P);

crates/bevy_ecs/hecs/src/lib.rs

+3-4
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ mod borrow;
6969
mod bundle;
7070
mod entities;
7171
mod entity_builder;
72+
mod filter;
7273
mod query;
7374
#[cfg(feature = "serde")]
7475
mod serde;
@@ -80,10 +81,8 @@ pub use borrow::{AtomicBorrow, Ref, RefMut};
8081
pub use bundle::{Bundle, DynamicBundle, MissingComponent};
8182
pub use entities::{Entity, EntityReserver, Location, NoSuchEntity};
8283
pub use entity_builder::{BuiltEntity, EntityBuilder};
83-
pub use query::{
84-
Added, Batch, BatchedIter, Changed, Mut, Mutated, Or, Query, QueryIter, ReadOnlyFetch, With,
85-
Without,
86-
};
84+
pub use filter::{Added, Changed, EntityFilter, Mutated, Or, QueryFilter, With, Without};
85+
pub use query::{Batch, BatchedIter, Mut, Query, QueryIter, ReadOnlyFetch};
8786
pub use world::{ArchetypesGeneration, Component, ComponentError, SpawnBatchIter, World};
8887

8988
// Unstable implementation details needed by the macros

0 commit comments

Comments
 (0)