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