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