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 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 get_mut(&self, entity_id: u32) -> Option<Q::Item<'_>> {
229 self.get(entity_id)
230 }
231
232 #[inline]
233 pub fn entity_count(&self) -> usize {
234 self.matching_archetypes
235 .iter()
236 .map(|&idx| self.world.archetype_index.archetypes[idx].len())
237 .sum()
238 }
239
240 #[inline]
241 pub fn len(&self) -> usize {
242 self.entity_count()
243 }
244
245 #[inline]
246 pub fn is_empty(&self) -> bool {
247 self.entity_count() == 0
248 }
249
250 #[inline]
252 pub fn contains(&self, entity_id: u32) -> bool {
253 self.get(entity_id).is_some()
254 }
255
256 pub fn entities<'a>(&'a self) -> impl Iterator<Item = u32> + 'a {
257 self.iter().map(|(id, _)| id)
258 }
259
260 pub fn par_for_each<F>(&self, func: F)
262 where
263 F: Fn((u32, Q::Item<'_>)) + Send + Sync,
264 {
265 use rayon::prelude::*;
266
267 #[derive(Copy, Clone)]
269 struct FetchWrapper<T>(T);
270 unsafe impl<T> Send for FetchWrapper<T> {}
271 unsafe impl<T> Sync for FetchWrapper<T> {}
272
273 impl<T: Copy> FetchWrapper<T> {
274 fn get(&self) -> T {
275 self.0
276 }
277 }
278
279 let tick = self.world.tick;
280 self.matching_archetypes.par_iter().for_each(|&arch_idx| {
281 let arch = &self.world.archetype_index.archetypes[arch_idx];
282 if let Some(fetch) = unsafe { Q::fetch_raw(arch, tick) } {
283 let len = arch.len();
284 let wrapped_fetch = FetchWrapper(fetch);
285 let entities_ptr = FetchWrapper(arch.entities().as_ptr());
286 let func_ref = &func;
287
288 (0..len)
291 .into_par_iter()
292 .with_min_len(512)
293 .for_each(move |row| unsafe {
294 if Q::filter_row(wrapped_fetch.get(), row, tick) {
295 let id = *entities_ptr.get().add(row);
296 let item = Q::get_item(wrapped_fetch.get(), row);
297 func_ref((id, item));
298 }
299 });
300 }
301 });
302 }
303
304 pub fn par_for_each_mut<F>(&mut self, func: F)
305 where
306 F: Fn((u32, Q::Item<'_>)) + Send + Sync,
307 {
308 self.par_for_each(func);
309 }
310}
311
312pub struct QueryIter<'a, 'w, Q: WorldQuery> {
317 world: &'a World,
318 archetype_indices: &'a [usize],
319 current_arch_idx: usize,
320 current_row: usize,
321 current_fetch: Option<Q::Fetch<'a>>,
322 _marker: PhantomData<Q>,
323 _marker_w: PhantomData<&'w ()>,
324}
325
326impl<'a, 'w, Q: WorldQuery> Iterator for QueryIter<'a, 'w, Q>
327where
328 'w: 'a,
329{
330 type Item = (u32, Q::Item<'a>);
331
332 fn next(&mut self) -> Option<Self::Item> {
333 loop {
334 if self.current_arch_idx >= self.archetype_indices.len() {
335 return None;
336 }
337
338 let arch_idx = self.archetype_indices[self.current_arch_idx];
339 let arch = &self.world.archetype_index.archetypes[arch_idx];
340
341 let fetch = match self.current_fetch {
342 Some(f) => f,
343 None => {
344 match unsafe { Q::fetch_raw(arch, self.world.tick) } {
345 Some(f) => {
346 self.current_fetch = Some(f);
347 self.current_row = 0;
348 f
349 }
350 None => {
351 self.current_arch_idx += 1;
353 continue;
354 }
355 }
356 }
357 };
358
359 if self.current_row < arch.len() {
360 let row = self.current_row;
361 self.current_row += 1;
362 if unsafe { Q::filter_row(fetch, row, self.world.tick) } {
363 let id = arch.entities()[row];
364 let item = unsafe { Q::get_item(fetch, row) };
365 return Some((id, item));
366 }
367 continue;
368 }
369
370 self.current_fetch = None;
371 self.current_arch_idx += 1;
372 }
373 }
374}
375
376pub struct QueryChunksIter<'a, 'w, Q: WorldQuery> {
381 world: &'a World,
382 archetype_indices: &'a [usize],
383 current_arch_idx: usize,
384 _marker: PhantomData<&'w Q>,
385}
386
387impl<'a, 'w, Q: WorldQuery> Iterator for QueryChunksIter<'a, 'w, Q>
388where
389 'w: 'a,
390{
391 type Item = (&'a [u32], Q::Slice<'a>);
392
393 fn next(&mut self) -> Option<Self::Item> {
394 while self.current_arch_idx < self.archetype_indices.len() {
395 let arch_idx = self.archetype_indices[self.current_arch_idx];
396 self.current_arch_idx += 1;
397
398 let arch = &self.world.archetype_index.archetypes[arch_idx];
399 let len = arch.len();
400 if len == 0 {
401 continue;
402 }
403
404 let fetch = match unsafe { Q::fetch_raw(arch, self.world.tick) } {
405 Some(f) => f,
406 None => continue,
407 };
408
409 let ids = unsafe { std::slice::from_raw_parts(arch.entities().as_ptr(), len) };
410 let slice = unsafe { Q::get_slice(fetch, len) };
411
412 return Some((ids, slice));
413 }
414 None
415 }
416}
417
418#[inline]
436fn check(tid: TypeId, is_mut: bool, types: &mut Vec<(TypeId, bool)>) {
437 for &(existing_tid, existing_mut) in types.iter() {
438 if existing_tid == tid && (existing_mut || is_mut) {
439 panic!(
440 "Query aliasing UB detected! Component TypeId {:?} is accessed mutably more than once \
441 in the same query. This would cause undefined behavior. \
442 Use separate queries for components of the same type that need independent mutable access.",
443 tid
444 );
445 }
446 }
447 types.push((tid, is_mut));
448}
449
450impl<T0: FetchComponent> WorldQuery for T0 {
451 type StaticType = T0::Component;
452 type Fetch<'w> = T0::Fetch<'w>;
453 type Item<'w> = T0::Item<'w>;
454 type Slice<'w> = T0::Slice<'w>;
455
456 unsafe fn fetch_raw<'w>(arch: &Archetype, tick: u32) -> Option<Self::Fetch<'w>> {
457 T0::fetch_raw(arch, tick)
458 }
459 fn check_aliasing(types: &mut Vec<(TypeId, bool)>) {
460 check(TypeId::of::<T0::Component>(), T0::IS_MUT, types);
461 }
462 fn matches_archetype(arch: &Archetype) -> bool {
463 arch.has_component(TypeId::of::<T0::Component>())
464 }
465
466 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w> {
467 T0::get_item(fetch, row)
468 }
469
470 unsafe fn filter_row<'w>(_fetch: Self::Fetch<'w>, _row: usize, _tick: u32) -> bool {
471 true
472 }
473
474 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w> {
475 T0::get_slice(fetch, len)
476 }
477}
478
479pub struct Changed<T>(PhantomData<T>);
480
481impl<T: 'static> WorldQuery for Changed<T> {
482 type StaticType = Changed<T>;
483 type Fetch<'w> = *const crate::archetype::ComponentTicks;
484 type Item<'w> = ();
485 type Slice<'w> = ();
486
487 unsafe fn fetch_raw<'w>(arch: &Archetype, _tick: u32) -> Option<Self::Fetch<'w>> {
488 let col = arch.get_column(TypeId::of::<T>())?;
489 Some(col.ticks_ptr())
490 }
491
492 fn check_aliasing(_types: &mut Vec<(TypeId, bool)>) {}
493
494 fn matches_archetype(arch: &Archetype) -> bool {
495 arch.has_component(TypeId::of::<T>())
496 }
497
498 unsafe fn filter_row<'w>(fetch: Self::Fetch<'w>, row: usize, tick: u32) -> bool {
499 let ticks = &*fetch.add(row);
500 ticks.changed == tick
501 }
502
503 unsafe fn get_item<'w>(_fetch: Self::Fetch<'w>, _row: usize) -> Self::Item<'w> {}
504
505 unsafe fn get_slice<'w>(_fetch: Self::Fetch<'w>, _len: usize) -> Self::Slice<'w> {}
506}
507
508macro_rules! impl_query_tuple {
509 ($($t:ident),*) => {
510 #[allow(non_snake_case)]
511 impl<$($t: WorldQuery),*> WorldQuery for ($($t,)*) {
512 type StaticType = ($($t::StaticType,)*);
513 type Fetch<'w> = ($($t::Fetch<'w>,)*);
514 type Item<'w> = ($($t::Item<'w>,)*);
515 type Slice<'w> = ($($t::Slice<'w>,)*);
516
517 unsafe fn fetch_raw<'w>(arch: &Archetype, tick: u32) -> Option<Self::Fetch<'w>> {
518 Some(($($t::fetch_raw(arch, tick)?,)*))
519 }
520 fn check_aliasing(types: &mut Vec<(TypeId, bool)>) {
521 $($t::check_aliasing(types);)*
522 }
523 fn matches_archetype(arch: &Archetype) -> bool {
524 $($t::matches_archetype(arch) &&)* true
525 }
526 unsafe fn get_item<'w>(fetch: Self::Fetch<'w>, row: usize) -> Self::Item<'w> {
527 let ($($t,)*) = fetch;
528 ($($t::get_item($t, row),)*)
529 }
530 unsafe fn filter_row<'w>(fetch: Self::Fetch<'w>, row: usize, tick: u32) -> bool {
531 let ($($t,)*) = fetch;
532 $($t::filter_row($t, row, tick) &&)* true
533 }
534 unsafe fn get_slice<'w>(fetch: Self::Fetch<'w>, len: usize) -> Self::Slice<'w> {
535 let ($($t,)*) = fetch;
536 ($($t::get_slice($t, len),)*)
537 }
538 }
539 };
540}
541
542impl_query_tuple!(T0, T1);
543impl_query_tuple!(T0, T1, T2);
544impl_query_tuple!(T0, T1, T2, T3);
545impl_query_tuple!(T0, T1, T2, T3, T4);
546impl_query_tuple!(T0, T1, T2, T3, T4, T5);
547impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6);
548impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7);
549impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8);
550impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9);
551impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
552impl_query_tuple!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
553
554pub struct With<T>(PhantomData<T>);
559
560impl<T: 'static> WorldQuery for With<T> {
561 type StaticType = With<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 Without<T>(PhantomData<T>);
584
585impl<T: 'static> WorldQuery for Without<T> {
586 type StaticType = Without<T>;
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 !arch.has_component(TypeId::of::<T>())
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
608pub struct Or<T1, T2>(PhantomData<(T1, T2)>);
609
610impl<T1: WorldQuery, T2: WorldQuery> WorldQuery for Or<T1, T2> {
611 type StaticType = Or<T1::StaticType, T2::StaticType>;
612 type Fetch<'w> = ();
613 type Item<'w> = ();
614 type Slice<'w> = ();
615
616 unsafe fn fetch_raw<'w>(_arch: &Archetype, _tick: u32) -> Option<Self::Fetch<'w>> {
617 Some(())
618 }
619
620 fn check_aliasing(_types: &mut Vec<(TypeId, bool)>) {}
621
622 fn matches_archetype(arch: &Archetype) -> bool {
623 T1::matches_archetype(arch) || T2::matches_archetype(arch)
624 }
625
626 unsafe fn filter_row<'w>(_fetch: Self::Fetch<'w>, _row: usize, _tick: u32) -> bool {
627 true
628 }
629 unsafe fn get_item<'w>(_fetch: Self::Fetch<'w>, _row: usize) -> Self::Item<'w> {}
630 unsafe fn get_slice<'w>(_fetch: Self::Fetch<'w>, _len: usize) -> Self::Slice<'w> {}
631}
632
633#[cfg(test)]
634mod tests {
635 use super::*;
636 use crate::impl_component;
637
638 #[derive(Debug, Clone, PartialEq)]
639 struct Position {
640 x: f32,
641 y: f32,
642 }
643 impl_component!(Position);
644
645 #[derive(Debug, Clone, PartialEq)]
646 struct Velocity {
647 x: f32,
648 y: f32,
649 }
650 impl_component!(Velocity);
651
652 #[test]
655 #[should_panic(expected = "Query aliasing UB detected")]
656 fn test_same_type_mut_mut_panics() {
657 let mut types = Vec::new();
658 check(TypeId::of::<Position>(), true, &mut types);
660 check(TypeId::of::<Position>(), true, &mut types);
662 }
663
664 #[test]
667 #[should_panic(expected = "Query aliasing UB detected")]
668 fn test_same_type_ref_mut_panics() {
669 let mut types = Vec::new();
670 check(TypeId::of::<Position>(), false, &mut types); check(TypeId::of::<Position>(), true, &mut types); }
673
674 #[test]
676 fn test_different_types_mut_mut_ok() {
677 let mut types = Vec::new();
678 check(TypeId::of::<Position>(), true, &mut types);
679 check(TypeId::of::<Velocity>(), true, &mut types);
680 assert_eq!(types.len(), 2);
681 }
682
683 #[test]
685 fn test_same_type_ref_ref_ok() {
686 let mut types = Vec::new();
687 check(TypeId::of::<Position>(), false, &mut types);
688 check(TypeId::of::<Position>(), false, &mut types);
689 assert_eq!(types.len(), 2);
690 }
691
692 #[test]
694 fn test_query_new_with_valid_types() {
695 let mut world = crate::World::new();
696 world.register_component_type::<Position>();
697 world.register_component_type::<Velocity>();
698 let e = world.spawn();
699 world.add_component(e, Position { x: 1.0, y: 2.0 });
700 world.add_component(e, Velocity { x: 0.0, y: 0.0 });
701
702 let q = world.query::<(Mut<Position>, Mut<Velocity>)>();
704 assert!(q.is_some());
705 }
706}