1use crate::archetype::Archetype;
2use crate::world::World;
3use std::any::TypeId;
4use std::marker::PhantomData;
5
6pub trait FetchComponent {
11 type Component: 'static;
12 type Fetch<'w>: Copy; type Item<'w>;
14 type Slice<'w>;
15
16 const IS_MUT: bool;
17
18 unsafe fn fetch_raw<'w>(arch: &Archetype, system_tick: u32) -> Option<Self::Fetch<'w>>;
23
24 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w>;
29
30 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w>;
35}
36
37impl<T: 'static> FetchComponent for &T {
38 type Component = T;
39 type Fetch<'w> = *const u8;
40 type Item<'w> = &'w T;
41 type Slice<'w> = &'w [T];
42 const IS_MUT: bool = false;
43
44 unsafe fn fetch_raw<'w>(arch: &Archetype, _system_tick: u32) -> Option<Self::Fetch<'w>> {
45 let col = arch.get_column(TypeId::of::<T>())?;
46 Some(col.data_ptr())
47 }
48
49 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w> {
50 let ptr = fetch.add(row * std::mem::size_of::<T>()) as *const T;
51 &*ptr
52 }
53
54 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w> {
55 std::slice::from_raw_parts(fetch as *const T, len)
56 }
57}
58
59pub struct Mut<'a, T: 'static> {
60 value: &'a mut T,
61 ticks: &'a mut crate::archetype::ComponentTicks,
62 current_tick: u32,
63}
64
65impl<T> std::ops::Deref for Mut<'_, T> {
66 type Target = T;
67 #[inline]
68 fn deref(&self) -> &T {
69 self.value
70 }
71}
72
73impl<T> std::ops::DerefMut for Mut<'_, T> {
74 #[inline]
75 fn deref_mut(&mut self) -> &mut T {
76 self.ticks.changed = self.current_tick;
77 self.value
78 }
79}
80
81impl<T: 'static> FetchComponent for Mut<'_, T> {
82 type Component = T;
83 type Fetch<'w> = (*mut u8, *mut crate::archetype::ComponentTicks, u32);
84 type Item<'w> = Mut<'w, T>;
85 type Slice<'w> = &'w mut [T];
86 const IS_MUT: bool = true;
87
88 unsafe fn fetch_raw<'w>(arch: &Archetype, system_tick: u32) -> Option<Self::Fetch<'w>> {
89 let mut col = arch.get_column_mut(TypeId::of::<T>())?;
90 Some((col.data_ptr_mut(), col.ticks_ptr_mut(), system_tick))
91 }
92
93 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w> {
94 let (data_ptr, ticks_ptr, system_tick) = fetch;
95 let ptr = data_ptr.add(row * std::mem::size_of::<T>()) as *mut T;
96 Mut {
97 value: &mut *ptr,
98 ticks: &mut *ticks_ptr.add(row),
99 current_tick: system_tick,
100 }
101 }
102
103 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w> {
104 let (data_ptr, ticks_ptr, system_tick) = fetch;
105
106 let ticks = std::slice::from_raw_parts_mut(ticks_ptr, len);
107 for tick in ticks.iter_mut() {
108 tick.changed = system_tick;
109 }
110
111 std::slice::from_raw_parts_mut(data_ptr as *mut T, len)
112 }
113}
114
115pub trait WorldQuery {
120 type StaticType: 'static;
121 type Fetch<'w>: Copy;
122 type Item<'w>;
123 type Slice<'w>;
124
125 unsafe fn fetch_raw<'w>(arch: &Archetype, system_tick: u32) -> Option<Self::Fetch<'w>>;
128 fn check_aliasing(types: &mut Vec<(TypeId, bool)>);
129 fn matches_archetype(arch: &Archetype) -> bool;
130
131 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w>;
134
135 unsafe fn filter_row<'w>(fetch: Self::Fetch<'w>, row: usize, system_tick: u32) -> bool;
138
139 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w>;
142}
143
144pub struct Query<'w, Q: WorldQuery + ?Sized> {
149 world: &'w World,
150 matching_archetypes: Vec<usize>,
151 _marker: PhantomData<Q>,
152}
153
154impl<'w, Q: WorldQuery> Query<'w, Q> {
155 pub fn new(world: &'w World) -> Option<Self> {
156 let mut used_types = Vec::new();
157 Q::check_aliasing(&mut used_types);
158 let matching = world
159 .archetype_index
160 .matching_archetypes_readonly(Q::matches_archetype);
161 Some(Self {
162 world,
163 matching_archetypes: matching,
164 _marker: PhantomData,
165 })
166 }
167
168 pub fn new_cached(world: &'w mut World) -> Option<Self> {
169 let mut used_types = Vec::new();
170 Q::check_aliasing(&mut used_types);
171 let matching = world
172 .archetype_index
173 .matching_archetypes(TypeId::of::<Q::StaticType>(), Q::matches_archetype)
174 .to_vec();
175 Some(Self {
176 world,
177 matching_archetypes: matching,
178 _marker: PhantomData,
179 })
180 }
181
182 pub fn iter<'a>(&'a self) -> QueryIter<'a, 'w, Q> {
183 QueryIter {
184 world: self.world,
185 archetype_indices: &self.matching_archetypes,
186 current_arch_idx: 0,
187 current_row: 0,
188 current_fetch: None,
189 _marker: PhantomData,
190 _marker_w: PhantomData,
191 }
192 }
193
194 pub fn iter_mut<'a>(&'a mut self) -> QueryIter<'a, 'w, Q> {
195 self.iter()
196 }
197
198 pub fn iter_chunks<'a>(&'a self) -> QueryChunksIter<'a, 'w, Q> {
199 QueryChunksIter {
200 world: self.world,
201 archetype_indices: &self.matching_archetypes,
202 current_arch_idx: 0,
203 _marker: PhantomData,
204 }
205 }
206
207 pub fn iter_chunks_mut<'a>(&'a mut self) -> QueryChunksIter<'a, 'w, Q> {
208 self.iter_chunks()
209 }
210
211 #[inline]
212 pub fn get(&self, entity_id: u32) -> Option<Q::Item<'_>> {
213 let loc = self.world.entity_location(entity_id);
214 if !loc.is_valid() {
215 return None;
216 }
217 let arch = &self.world.archetype_index.archetypes[loc.archetype_id as usize];
218 unsafe {
219 let fetch = Q::fetch_raw(arch, self.world.tick)?;
220 if !Q::filter_row(fetch, loc.row as usize, self.world.tick) {
221 return None;
222 }
223 Some(Q::get_item(fetch, loc.row as usize))
224 }
225 }
226
227 #[inline]
228 pub fn entity_count(&self) -> usize {
229 self.matching_archetypes
230 .iter()
231 .map(|&idx| self.world.archetype_index.archetypes[idx].len())
232 .sum()
233 }
234
235 pub fn par_for_each<F>(&self, func: F)
237 where
238 F: Fn((u32, Q::Item<'_>)) + Send + Sync,
239 {
240 use rayon::prelude::*;
241
242 #[derive(Copy, Clone)]
244 struct FetchWrapper<T>(T);
245 unsafe impl<T> Send for FetchWrapper<T> {}
246 unsafe impl<T> Sync for FetchWrapper<T> {}
247
248 impl<T: Copy> FetchWrapper<T> {
249 fn get(&self) -> T {
250 self.0
251 }
252 }
253
254 let tick = self.world.tick;
255 self.matching_archetypes.par_iter().for_each(|&arch_idx| {
256 let arch = &self.world.archetype_index.archetypes[arch_idx];
257 if let Some(fetch) = unsafe { Q::fetch_raw(arch, tick) } {
258 let len = arch.len();
259 let wrapped_fetch = FetchWrapper(fetch);
260 let entities_ptr = FetchWrapper(arch.entities().as_ptr());
261 let func_ref = &func;
262
263 (0..len)
266 .into_par_iter()
267 .with_min_len(512)
268 .for_each(move |row| unsafe {
269 if Q::filter_row(wrapped_fetch.get(), row, tick) {
270 let id = *entities_ptr.get().add(row);
271 let item = Q::get_item(wrapped_fetch.get(), row);
272 func_ref((id, item));
273 }
274 });
275 }
276 });
277 }
278
279 pub fn par_for_each_mut<F>(&mut self, func: F)
280 where
281 F: Fn((u32, Q::Item<'_>)) + Send + Sync,
282 {
283 self.par_for_each(func);
284 }
285}
286
287pub struct QueryIter<'a, 'w, Q: WorldQuery> {
292 world: &'a World,
293 archetype_indices: &'a [usize],
294 current_arch_idx: usize,
295 current_row: usize,
296 current_fetch: Option<Q::Fetch<'a>>,
297 _marker: PhantomData<Q>,
298 _marker_w: PhantomData<&'w ()>,
299}
300
301impl<'a, 'w, Q: WorldQuery> Iterator for QueryIter<'a, 'w, Q>
302where
303 'w: 'a,
304{
305 type Item = (u32, Q::Item<'a>);
306
307 fn next(&mut self) -> Option<Self::Item> {
308 loop {
309 if self.current_arch_idx >= self.archetype_indices.len() {
310 return None;
311 }
312
313 let arch_idx = self.archetype_indices[self.current_arch_idx];
314 let arch = &self.world.archetype_index.archetypes[arch_idx];
315
316 let fetch = match self.current_fetch {
317 Some(f) => f,
318 None => {
319 match unsafe { Q::fetch_raw(arch, self.world.tick) } {
320 Some(f) => {
321 self.current_fetch = Some(f);
322 self.current_row = 0;
323 f
324 }
325 None => {
326 self.current_arch_idx += 1;
328 continue;
329 }
330 }
331 }
332 };
333
334 if self.current_row < arch.len() {
335 let row = self.current_row;
336 self.current_row += 1;
337 if unsafe { Q::filter_row(fetch, row, self.world.tick) } {
338 let id = arch.entities()[row];
339 let item = unsafe { Q::get_item(fetch, row) };
340 return Some((id, item));
341 }
342 continue;
343 }
344
345 self.current_fetch = None;
346 self.current_arch_idx += 1;
347 }
348 }
349}
350
351pub struct QueryChunksIter<'a, 'w, Q: WorldQuery> {
356 world: &'a World,
357 archetype_indices: &'a [usize],
358 current_arch_idx: usize,
359 _marker: PhantomData<&'w Q>,
360}
361
362impl<'a, 'w, Q: WorldQuery> Iterator for QueryChunksIter<'a, 'w, Q>
363where
364 'w: 'a,
365{
366 type Item = (&'a [u32], Q::Slice<'a>);
367
368 fn next(&mut self) -> Option<Self::Item> {
369 while self.current_arch_idx < self.archetype_indices.len() {
370 let arch_idx = self.archetype_indices[self.current_arch_idx];
371 self.current_arch_idx += 1;
372
373 let arch = &self.world.archetype_index.archetypes[arch_idx];
374 let len = arch.len();
375 if len == 0 {
376 continue;
377 }
378
379 let fetch = match unsafe { Q::fetch_raw(arch, self.world.tick) } {
380 Some(f) => f,
381 None => continue,
382 };
383
384 let ids = unsafe { std::slice::from_raw_parts(arch.entities().as_ptr(), len) };
385 let slice = unsafe { Q::get_slice(fetch, len) };
386
387 return Some((ids, slice));
388 }
389 None
390 }
391}
392
393#[inline]
411fn check(tid: TypeId, is_mut: bool, types: &mut Vec<(TypeId, bool)>) {
412 for &(existing_tid, existing_mut) in types.iter() {
413 if existing_tid == tid && (existing_mut || is_mut) {
414 panic!(
415 "Query aliasing UB detected! Component TypeId {:?} is accessed mutably more than once \
416 in the same query. This would cause undefined behavior. \
417 Use separate queries for components of the same type that need independent mutable access.",
418 tid
419 );
420 }
421 }
422 types.push((tid, is_mut));
423}
424
425impl<T0: FetchComponent> WorldQuery for T0 {
426 type StaticType = T0::Component;
427 type Fetch<'w> = T0::Fetch<'w>;
428 type Item<'w> = T0::Item<'w>;
429 type Slice<'w> = T0::Slice<'w>;
430
431 unsafe fn fetch_raw<'w>(arch: &Archetype, tick: u32) -> Option<Self::Fetch<'w>> {
432 T0::fetch_raw(arch, tick)
433 }
434 fn check_aliasing(types: &mut Vec<(TypeId, bool)>) {
435 check(TypeId::of::<T0::Component>(), T0::IS_MUT, types);
436 }
437 fn matches_archetype(arch: &Archetype) -> bool {
438 arch.has_component(TypeId::of::<T0::Component>())
439 }
440
441 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w> {
442 T0::get_item(fetch, row)
443 }
444
445 unsafe fn filter_row<'w>(_fetch: Self::Fetch<'w>, _row: usize, _tick: u32) -> bool {
446 true
447 }
448
449 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w> {
450 T0::get_slice(fetch, len)
451 }
452}
453
454pub struct Changed<T>(PhantomData<T>);
455
456impl<T: 'static> WorldQuery for Changed<T> {
457 type StaticType = Changed<T>;
458 type Fetch<'w> = *const crate::archetype::ComponentTicks;
459 type Item<'w> = ();
460 type Slice<'w> = ();
461
462 unsafe fn fetch_raw<'w>(arch: &Archetype, _tick: u32) -> Option<Self::Fetch<'w>> {
463 let col = arch.get_column(TypeId::of::<T>())?;
464 Some(col.ticks_ptr())
465 }
466
467 fn check_aliasing(_types: &mut Vec<(TypeId, bool)>) {}
468
469 fn matches_archetype(arch: &Archetype) -> bool {
470 arch.has_component(TypeId::of::<T>())
471 }
472
473 unsafe fn filter_row<'w>(fetch: Self::Fetch<'w>, row: usize, tick: u32) -> bool {
474 let ticks = &*fetch.add(row);
475 ticks.changed == tick
476 }
477
478 unsafe fn get_item<'w>(_fetch: Self::Fetch<'w>, _row: usize) -> Self::Item<'w> {}
479
480 unsafe fn get_slice<'w>(_fetch: Self::Fetch<'w>, _len: usize) -> Self::Slice<'w> {}
481}
482
483macro_rules! impl_query_tuple {
484 ($($t:ident),*) => {
485 #[allow(non_snake_case)]
486 impl<$($t: WorldQuery),*> WorldQuery for ($($t,)*) {
487 type StaticType = ($($t::StaticType,)*);
488 type Fetch<'w> = ($($t::Fetch<'w>,)*);
489 type Item<'w> = ($($t::Item<'w>,)*);
490 type Slice<'w> = ($($t::Slice<'w>,)*);
491
492 unsafe fn fetch_raw<'w>(arch: &Archetype, tick: u32) -> Option<Self::Fetch<'w>> {
493 Some(($($t::fetch_raw(arch, tick)?,)*))
494 }
495 fn check_aliasing(types: &mut Vec<(TypeId, bool)>) {
496 $($t::check_aliasing(types);)*
497 }
498 fn matches_archetype(arch: &Archetype) -> bool {
499 $($t::matches_archetype(arch) &&)* true
500 }
501 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w> {
502 let ($($t,)*) = fetch;
503 ($($t::get_item($t, row),)*)
504 }
505 unsafe fn filter_row<'w>(fetch: Self::Fetch<'w>, row: usize, tick: u32) -> bool {
506 let ($($t,)*) = fetch;
507 $($t::filter_row($t, row, tick) &&)* true
508 }
509 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w> {
510 let ($($t,)*) = fetch;
511 ($($t::get_slice($t, len),)*)
512 }
513 }
514 };
515}
516
517impl_query_tuple!(T0, T1);
518impl_query_tuple!(T0, T1, T2);
519impl_query_tuple!(T0, T1, T2, T3);
520impl_query_tuple!(T0, T1, T2, T3, T4);
521impl_query_tuple!(T0, T1, T2, T3, T4, T5);
522impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6);
523impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7);
524impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8);
525impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9);
526impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
527impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
528
529pub struct With<T>(PhantomData<T>);
534
535impl<T: 'static> WorldQuery for With<T> {
536 type StaticType = With<T>;
537 type Fetch<'w> = ();
538 type Item<'w> = ();
539 type Slice<'w> = ();
540
541 unsafe fn fetch_raw<'w>(_arch: &Archetype, _tick: u32) -> Option<Self::Fetch<'w>> {
542 Some(())
543 }
544
545 fn check_aliasing(_types: &mut Vec<(TypeId, bool)>) {}
546
547 fn matches_archetype(arch: &Archetype) -> bool {
548 arch.has_component(TypeId::of::<T>())
549 }
550
551 unsafe fn filter_row<'w>(_fetch: Self::Fetch<'w>, _row: usize, _tick: u32) -> bool {
552 true
553 }
554 unsafe fn get_item<'w>(_fetch: Self::Fetch<'w>, _row: usize) -> Self::Item<'w> {}
555 unsafe fn get_slice<'w>(_fetch: Self::Fetch<'w>, _len: usize) -> Self::Slice<'w> {}
556}
557
558pub struct Without<T>(PhantomData<T>);
559
560impl<T: 'static> WorldQuery for Without<T> {
561 type StaticType = Without<T>;
562 type Fetch<'w> = ();
563 type Item<'w> = ();
564 type Slice<'w> = ();
565
566 unsafe fn fetch_raw<'w>(_arch: &Archetype, _tick: u32) -> Option<Self::Fetch<'w>> {
567 Some(())
568 }
569
570 fn check_aliasing(_types: &mut Vec<(TypeId, bool)>) {}
571
572 fn matches_archetype(arch: &Archetype) -> bool {
573 !arch.has_component(TypeId::of::<T>())
574 }
575
576 unsafe fn filter_row<'w>(_fetch: Self::Fetch<'w>, _row: usize, _tick: u32) -> bool {
577 true
578 }
579 unsafe fn get_item<'w>(_fetch: Self::Fetch<'w>, _row: usize) -> Self::Item<'w> {}
580 unsafe fn get_slice<'w>(_fetch: Self::Fetch<'w>, _len: usize) -> Self::Slice<'w> {}
581}
582
583pub struct Or<T1, T2>(PhantomData<(T1, T2)>);
584
585impl<T1: WorldQuery, T2: WorldQuery> WorldQuery for Or<T1, T2> {
586 type StaticType = Or<T1::StaticType, T2::StaticType>;
587 type Fetch<'w> = ();
588 type Item<'w> = ();
589 type Slice<'w> = ();
590
591 unsafe fn fetch_raw<'w>(_arch: &Archetype, _tick: u32) -> Option<Self::Fetch<'w>> {
592 Some(())
593 }
594
595 fn check_aliasing(_types: &mut Vec<(TypeId, bool)>) {}
596
597 fn matches_archetype(arch: &Archetype) -> bool {
598 T1::matches_archetype(arch) || T2::matches_archetype(arch)
599 }
600
601 unsafe fn filter_row<'w>(_fetch: Self::Fetch<'w>, _row: usize, _tick: u32) -> bool {
602 true
603 }
604 unsafe fn get_item<'w>(_fetch: Self::Fetch<'w>, _row: usize) -> Self::Item<'w> {}
605 unsafe fn get_slice<'w>(_fetch: Self::Fetch<'w>, _len: usize) -> Self::Slice<'w> {}
606}
607
608#[cfg(test)]
609mod tests {
610 use super::*;
611 use crate::impl_component;
612
613 #[derive(Debug, Clone, PartialEq)]
614 struct Position {
615 x: f32,
616 y: f32,
617 }
618 impl_component!(Position);
619
620 #[derive(Debug, Clone, PartialEq)]
621 struct Velocity {
622 x: f32,
623 y: f32,
624 }
625 impl_component!(Velocity);
626
627 #[test]
630 #[should_panic(expected = "Query aliasing UB detected")]
631 fn test_same_type_mut_mut_panics() {
632 let mut types = Vec::new();
633 check(TypeId::of::<Position>(), true, &mut types);
635 check(TypeId::of::<Position>(), true, &mut types);
637 }
638
639 #[test]
642 #[should_panic(expected = "Query aliasing UB detected")]
643 fn test_same_type_ref_mut_panics() {
644 let mut types = Vec::new();
645 check(TypeId::of::<Position>(), false, &mut types); check(TypeId::of::<Position>(), true, &mut types); }
648
649 #[test]
651 fn test_different_types_mut_mut_ok() {
652 let mut types = Vec::new();
653 check(TypeId::of::<Position>(), true, &mut types);
654 check(TypeId::of::<Velocity>(), true, &mut types);
655 assert_eq!(types.len(), 2);
656 }
657
658 #[test]
660 fn test_same_type_ref_ref_ok() {
661 let mut types = Vec::new();
662 check(TypeId::of::<Position>(), false, &mut types);
663 check(TypeId::of::<Position>(), false, &mut types);
664 assert_eq!(types.len(), 2);
665 }
666
667 #[test]
669 fn test_query_new_with_valid_types() {
670 let mut world = crate::World::new();
671 world.register_component_type::<Position>();
672 world.register_component_type::<Velocity>();
673 let e = world.spawn();
674 world.add_component(e, Position { x: 1.0, y: 2.0 });
675 world.add_component(e, Velocity { x: 0.0, y: 0.0 });
676
677 let q = world.query::<(Mut<Position>, Mut<Velocity>)>();
679 assert!(q.is_some());
680 }
681}