rs_ecs/
query.rs

1use std::any::TypeId;
2use std::iter::FusedIterator;
3use std::marker::PhantomData;
4use std::mem::{transmute, transmute_copy, MaybeUninit};
5use std::ptr::NonNull;
6use std::slice::Iter;
7
8use crate::{
9    archetype::Archetype,
10    borrow_flags::{BorrowFlags, Ref, RefMut},
11    world::{Entity, EntityMetadata, World},
12};
13
14#[cfg(feature = "rayon")]
15use crate::rayon::QueryParIter;
16
17/// Query to get an iterator over all entities with a certain combination of components.
18///
19/// Queries are specified through their type argument, by composing the type of their result.
20/// The following type specifications are possible:
21///
22/// * `&C` - shared, immutable reference to components of type `C`
23/// * `&mut C` - exclusive, mutable reference
24/// * `(P, Q, R)` - combine multiple queries
25/// * `Option<Q>` - optional component(s)
26/// * `With<Q, C>` to filter `Q` for presence of `C`
27/// * `Without<Q, C>` to filter `Q` for absence of `C`
28/// * `Matches<Q> to indicate which entities match `Q`
29///
30/// Note that [Entities](Entity) are components themselves, so they can be optionally obtained in a query,
31/// like `Query<Entity, &C, &mut D>`.
32///
33/// Queries are provided as stand-alone structs to allow for prepared queries that can be
34/// re-used, as an optimzation. Hence, queries need to borrow the [World] before their results
35/// can be iterated (see [Query::borrow]).
36///
37/// # Examples
38///
39/// ```
40/// # use rs_ecs::*;
41/// let mut world = World::new();
42///
43/// let entity1 = world.alloc();
44/// world.insert(entity1, (42_i32, 23_u32, 1.0_f32));
45///
46/// let entity2 = world.alloc();
47/// world.insert(entity2, (0_i32, true));
48///
49/// let mut query = Query::<(&mut i32, Option<(&u32, &mut f32)>)>::new();
50/// for (i, opt) in query.borrow(&world).iter() {
51///     if let Some((u, f)) = opt {
52///         *i += *u as i32;
53///         *f += 1.0
54///     } else {
55///         *i -= 1;
56///     }
57/// }
58///
59/// assert_eq!(*world.query_one::<&i32>(entity1).unwrap().get(), 65);
60/// assert_eq!(*world.query_one::<&f32>(entity1).unwrap().get(), 2.0);
61/// assert_eq!(*world.query_one::<&i32>(entity2).unwrap().get(), -1);
62/// ```
63///
64/// Use of a prepared query that is stored and reused for optimization:
65///
66/// ```
67/// # use rs_ecs::*;
68/// #[derive(Default)]
69/// struct System {
70///     query: Query<(&'static i32, &'static mut bool)>
71/// }
72///
73/// impl System {
74///     pub fn update(&mut self, world: &World) {
75///         for (i, b) in self.query.borrow(world).iter() {
76///             *b = *i >= 0;
77///         }
78///     }
79/// }
80///
81/// fn main() {
82///     let mut world = World::new();
83///
84///     let entity = world.alloc();
85///     world.insert(entity, (23_i32, false));
86///
87///     let mut system = System::default();
88///
89///     for _ in 0..42 {
90///         system.update(&world);
91///     }
92/// }
93/// ```
94///
95pub struct Query<S>
96where
97    S: QuerySpec,
98{
99    tag_gen: (u32, u16),
100    flags: Option<<S::Fetch as Fetch<'static>>::Ty>,
101    comps: Box<[(u16, <S::Fetch as Fetch<'static>>::Ty)]>,
102    ptrs: Box<[Option<<S::Fetch as Fetch<'static>>::Ptr>]>,
103}
104
105unsafe impl<S> Send for Query<S> where S: QuerySpec {}
106
107impl<S> Default for Query<S>
108where
109    S: QuerySpec,
110{
111    /// Create a query.
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117impl<S> Query<S>
118where
119    S: QuerySpec,
120{
121    /// Create a query.
122    ///
123    /// # Examples
124    ///
125    /// ```
126    /// # use rs_ecs::*;
127    /// let mut immutable_query = Query::<(&i32,)>::new();
128    /// let mut mutable_query = Query::<(&i32, &mut bool)>::new();
129    /// let mut query_with_entity = Query::<(&Entity, &i32, &mut bool)>::new();
130    /// ```
131    pub fn new() -> Self {
132        Self {
133            tag_gen: Default::default(),
134            flags: Default::default(),
135            comps: Default::default(),
136            ptrs: Default::default(),
137        }
138    }
139
140    /// Borrow the [World] to allow for iterating the query.
141    ///
142    /// # Examples
143    ///
144    /// ```
145    /// # use rs_ecs::*;
146    /// let mut world = World::new();
147    /// let mut query = Query::<(&i32, &bool)>::new();
148    /// let mut borrow = query.borrow(&world);
149    ///
150    /// for (i, b) in borrow.iter() {
151    ///     println!("{}, {}", i, b);
152    /// }
153    /// ```
154    pub fn borrow<'w>(&'w mut self, world: &'w World) -> QueryRef<'w, S> {
155        if self.tag_gen != world.tag_gen() {
156            self.find(world);
157        }
158
159        let _ref = self
160            .flags
161            .map(|ty| unsafe { S::Fetch::borrow(&world.borrow_flags, transmute_copy(&ty)) });
162
163        QueryRef {
164            world,
165            _ref,
166            comps: unsafe { transmute(&*self.comps) },
167            ptrs: unsafe { transmute(&mut *self.ptrs) },
168        }
169    }
170
171    #[cold]
172    #[inline(never)]
173    fn find(&mut self, world: &World) {
174        self.flags = S::Fetch::find_flags(&world.borrow_flags);
175
176        self.comps = world
177            .archetypes
178            .iter()
179            .enumerate()
180            .filter_map(|(idx, archetype)| {
181                S::Fetch::find_comps(archetype).map(|ty| (idx as u16, ty))
182            })
183            .collect();
184
185        self.ptrs = world.archetypes.iter().map(|_| None).collect();
186
187        self.tag_gen = world.tag_gen();
188    }
189}
190
191/// Borrow of the [World] for a [Query]. Required to obtain an iterator.
192pub struct QueryRef<'w, S>
193where
194    S: QuerySpec,
195{
196    world: &'w World,
197    _ref: Option<<S::Fetch as Fetch<'w>>::Ref>,
198    comps: &'w [(u16, <S::Fetch as Fetch<'w>>::Ty)],
199    ptrs: &'w mut [Option<<S::Fetch as Fetch<'w>>::Ptr>],
200}
201
202impl<S> QueryRef<'_, S>
203where
204    S: QuerySpec,
205{
206    /// Visit all entities matching the query.
207    pub fn for_each<'q, F>(&'q mut self, mut f: F)
208    where
209        F: FnMut(<S::Fetch as Fetch<'q>>::Item),
210    {
211        let comps: &'q [(u16, <S::Fetch as Fetch<'q>>::Ty)] = unsafe { transmute(self.comps) };
212
213        for (idx, ty) in comps {
214            let archetype = &self.world.archetypes[*idx as usize];
215
216            let ptr = unsafe { S::Fetch::base_pointer(archetype, *ty) };
217
218            for idx in 0..archetype.len() {
219                let val = unsafe { S::Fetch::deref(ptr, idx) };
220
221                f(val);
222            }
223        }
224    }
225
226    /// Create an iterator over the entities matching the query.
227    pub fn iter<'q>(&'q mut self) -> QueryIter<'q, S> {
228        let comps: &'q [(u16, <S::Fetch as Fetch<'q>>::Ty)] = unsafe { transmute(self.comps) };
229
230        QueryIter {
231            comps: comps.iter(),
232            archetypes: &self.world.archetypes,
233            idx: 0,
234            len: 0,
235            ptr: S::Fetch::dangling(),
236        }
237    }
238
239    #[cfg(feature = "rayon")]
240    /// Create a parallel iterator over the entities matching the query.
241    ///
242    /// # Examples
243    ///
244    /// ```
245    /// # use rs_ecs::*;
246    /// # use rayon::prelude::*;
247    /// use rayon::join;
248    ///
249    /// let mut world = World::new();
250    ///
251    /// let entity1 = world.alloc();
252    /// world.insert(entity1, (42_i32, 23_u32, 1.0_f32));
253    ///
254    /// let entity2 = world.alloc();
255    /// world.insert(entity2, (0_i32, true));
256    ///
257    /// let mut query1 = Query::<&mut i32>::new();
258    /// let mut ref1 = query1.borrow(&world);
259    /// let mut iter1 = ref1.iter();
260    ///
261    /// let mut query2 = Query::<(&u32, &mut f32)>::new();
262    /// let mut ref2 = query2.borrow(&world);
263    /// let mut iter2 = ref2.par_iter();
264    ///
265    /// join(
266    ///     move || {
267    ///         for i in iter1 {
268    ///             *i -= 1;
269    ///         }
270    ///     },
271    ///     move || {
272    ///         iter2.for_each(|(u, f)| {
273    ///             *f += *u as f32;
274    ///         });
275    ///     },
276    /// );
277    /// ```
278    pub fn par_iter<'q>(&'q mut self) -> QueryParIter<'q, S>
279    where
280        <S::Fetch as Fetch<'q>>::Item: Send,
281    {
282        let comps: &'q [(u16, <S::Fetch as Fetch<'q>>::Ty)] = unsafe { transmute(self.comps) };
283
284        QueryParIter::new(comps, &self.world.archetypes)
285    }
286
287    /// Create a map of the entities matching the query.
288    ///
289    /// This is an alternative to [`query_one`](World::query_one) for repeated calls, to amortize the set-up costs of random access.
290    ///
291    /// # Examples
292    ///
293    /// ```
294    /// # use rs_ecs::*;
295    /// let mut world = World::new();
296    ///
297    /// let entity = world.alloc();
298    /// world.insert(entity, (42_i32, 1.0_f32));
299    ///
300    /// let mut query = Query::<(&i32, &f32)>::new();
301    /// let mut query = query.borrow(&world);
302    /// let mut query = query.map();
303    ///
304    /// let (i, f) = query.get_mut(entity).unwrap();
305    ///
306    /// assert_eq!(*i, 42);
307    /// assert_eq!(*f, 1.0);
308    /// ```
309    pub fn map<'q>(&'q mut self) -> QueryMap<'q, S> {
310        let comps: &'q [(u16, <S::Fetch as Fetch<'q>>::Ty)] = unsafe { transmute(self.comps) };
311
312        let ptrs: &'q mut [Option<<S::Fetch as Fetch<'q>>::Ptr>] =
313            unsafe { transmute(&mut *self.ptrs) };
314
315        ptrs.fill(None);
316
317        for (idx, ty) in comps {
318            let archetype = &self.world.archetypes[*idx as usize];
319
320            let ptr = unsafe { S::Fetch::base_pointer(archetype, *ty) };
321
322            ptrs[*idx as usize] = Some(ptr);
323        }
324
325        QueryMap {
326            entities: &self.world.entities,
327            ptrs,
328        }
329    }
330}
331
332/// Used to iterate through the entities which match a certain [Query].
333pub struct QueryIter<'q, S>
334where
335    S: QuerySpec,
336{
337    comps: Iter<'q, (u16, <S::Fetch as Fetch<'q>>::Ty)>,
338    archetypes: &'q [Archetype],
339    idx: u32,
340    len: u32,
341    ptr: <S::Fetch as Fetch<'q>>::Ptr,
342}
343
344unsafe impl<'q, S> Send for QueryIter<'q, S>
345where
346    S: QuerySpec,
347    <S::Fetch as Fetch<'q>>::Item: Send,
348{
349}
350
351impl<'q, S> Iterator for QueryIter<'q, S>
352where
353    S: QuerySpec,
354{
355    type Item = <S::Fetch as Fetch<'q>>::Item;
356
357    fn next(&mut self) -> Option<Self::Item> {
358        loop {
359            if self.idx != self.len {
360                let val = unsafe { S::Fetch::deref(self.ptr, self.idx) };
361                self.idx += 1;
362                return Some(val);
363            } else {
364                let (idx, ty) = self.comps.next()?;
365                let archetype = &self.archetypes[*idx as usize];
366                self.idx = 0;
367                self.len = archetype.len();
368                self.ptr = unsafe { S::Fetch::base_pointer(archetype, *ty) };
369            }
370        }
371    }
372
373    fn size_hint(&self) -> (usize, Option<usize>) {
374        let len = self.len();
375        (len, Some(len))
376    }
377}
378
379impl<S> ExactSizeIterator for QueryIter<'_, S>
380where
381    S: QuerySpec,
382{
383    fn len(&self) -> usize {
384        let len = self
385            .comps
386            .clone()
387            .map(|(idx, _)| self.archetypes[*idx as usize].len())
388            .sum::<u32>()
389            + self.len
390            - self.idx;
391        len as usize
392    }
393}
394
395impl<S> FusedIterator for QueryIter<'_, S> where S: QuerySpec {}
396
397/// Provides random access to the entities which match a certain [Query].
398pub struct QueryMap<'q, S>
399where
400    S: QuerySpec,
401{
402    entities: &'q [EntityMetadata],
403    ptrs: &'q [Option<<S::Fetch as Fetch<'q>>::Ptr>],
404}
405
406unsafe impl<'q, S> Send for QueryMap<'q, S>
407where
408    S: QuerySpec,
409    <S::Fetch as Fetch<'q>>::Item: Send,
410{
411}
412
413unsafe impl<'q, S> Sync for QueryMap<'q, S>
414where
415    S: QuerySpec,
416    <S::Fetch as Fetch<'q>>::Item: Send,
417{
418}
419
420impl<S> QueryMap<'_, S>
421where
422    S: QuerySpec,
423{
424    /// Check if an [Entity] is present in the map.
425    ///
426    /// In contrast to [`get`][Self::get], this does not require a `S::Fetch: FetchShared` bound.
427    /// In contrast to [`get_mut`][Self::get_mut] however, a shared borrow is sufficient.
428    ///
429    /// # Examples
430    ///
431    /// ```
432    /// # use rs_ecs::*;
433    /// let mut world = World::new();
434    ///
435    /// let entity1 = world.alloc();
436    /// world.insert(entity1, (42_i32, 1.0_f32));
437    ///
438    /// let entity2 = world.alloc();
439    /// world.insert(entity2, (2.0_f32,));
440    ///
441    /// let mut query = Query::<(&i32, Option<&f32>)>::new();
442    /// let mut query = query.borrow(&world);
443    /// let query = query.map();
444    ///
445    /// assert!(query.contains(entity1));
446    /// assert!(!query.contains(entity2));
447    /// ```
448    pub fn contains(&self, ent: Entity) -> bool {
449        let meta = self.entities[ent.id as usize];
450        assert_eq!(ent.gen, meta.gen, "Entity is stale");
451
452        let ptr = unsafe { self.ptrs.get_unchecked(meta.ty as usize) };
453
454        ptr.is_some()
455    }
456
457    /// Access the queried components of the given [Entity]
458    ///
459    /// Available only if the components do not include unique references.
460    ///
461    /// # Examples
462    ///
463    /// ```
464    /// # use rs_ecs::*;
465    /// let mut world = World::new();
466    ///
467    /// let entity1 = world.alloc();
468    /// world.insert(entity1, (42_i32, 1.0_f32));
469    ///
470    /// let entity2 = world.alloc();
471    /// world.insert(entity2, (23_i32,));
472    ///
473    /// let mut query = Query::<(&i32, Option<&f32>)>::new();
474    /// let mut query = query.borrow(&world);
475    /// let mut query = query.map();
476    ///
477    /// let (i1, f1) = query.get(entity1).unwrap();
478    /// let (i2, f2) = query.get(entity2).unwrap();
479    ///
480    /// assert_eq!(*i1, 42);
481    /// assert_eq!(f1.copied(), Some(1.0));
482    /// assert_eq!(*i2, 23);
483    /// assert_eq!(f2.copied(), None);
484    /// ```
485    pub fn get(&self, ent: Entity) -> Option<<S::Fetch as Fetch>::Item>
486    where
487        S::Fetch: FetchShared,
488    {
489        unsafe { self.get_unchecked(ent) }
490    }
491
492    /// Exclusively access the queried components of the given [Entity]
493    pub fn get_mut(&mut self, ent: Entity) -> Option<<S::Fetch as Fetch>::Item> {
494        unsafe { self.get_unchecked(ent) }
495    }
496
497    /// Exclusively access the queried component of the given distinct entities
498    ///
499    /// # Examples
500    ///
501    /// ```
502    /// # use rs_ecs::*;
503    /// let mut world = World::new();
504    ///
505    /// let entity1 = world.alloc();
506    /// world.insert(entity1, (42_i32, 1.0_f32));
507    ///
508    /// let entity2 = world.alloc();
509    /// world.insert(entity2, (23_i32,));
510    ///
511    /// let mut query = Query::<&mut i32>::new();
512    /// let mut query = query.borrow(&world);
513    /// let mut query = query.map();
514    ///
515    /// let [i1, i2] = query.get_many_mut([entity1, entity2]);
516    ///
517    /// assert_eq!(*i1.unwrap(), 42);
518    /// assert_eq!(*i2.unwrap(), 23);
519    /// ```
520    pub fn get_many_mut<const N: usize>(
521        &mut self,
522        ent: [Entity; N],
523    ) -> [Option<<S::Fetch as Fetch>::Item>; N] {
524        let mut val = MaybeUninit::uninit();
525
526        let ptr = val.as_mut_ptr() as *mut Option<<S::Fetch as Fetch>::Item>;
527
528        for idx in 0..N {
529            for idx1 in (idx + 1)..N {
530                assert_ne!(ent[idx], ent[idx1], "Duplicate entity");
531            }
532
533            unsafe {
534                let val = self.get_unchecked(ent[idx]);
535
536                ptr.add(idx).write(val);
537            }
538        }
539
540        unsafe { val.assume_init() }
541    }
542
543    unsafe fn get_unchecked<'m>(&'m self, ent: Entity) -> Option<<S::Fetch as Fetch<'m>>::Item> {
544        let meta = self.entities[ent.id as usize];
545        assert_eq!(ent.gen, meta.gen, "Entity is stale");
546
547        let ptr: &'m Option<<S::Fetch as Fetch<'m>>::Ptr> =
548            transmute(self.ptrs.get_unchecked(meta.ty as usize));
549
550        ptr.map(|ptr| S::Fetch::deref(ptr, meta.idx))
551    }
552}
553
554/// Type level specification of a query for a certain set of components.
555pub trait QuerySpec {
556    #[doc(hidden)]
557    type Fetch: for<'a> Fetch<'a>;
558}
559
560#[allow(clippy::missing_safety_doc)]
561pub unsafe trait Fetch<'q> {
562    type Ty: Copy + 'static;
563    type Ref;
564    type Ptr: Copy + 'static;
565
566    type Item;
567
568    fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty>;
569    fn find_comps(archetype: &Archetype) -> Option<Self::Ty>;
570    unsafe fn borrow(borrows: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref;
571    unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr;
572
573    fn dangling() -> Self::Ptr;
574    unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item;
575}
576
577#[allow(clippy::missing_safety_doc)]
578pub unsafe trait FetchShared {}
579
580impl<'a, C> QuerySpec for &'a C
581where
582    C: 'static,
583{
584    type Fetch = FetchRead<C>;
585}
586
587pub struct FetchRead<C>(PhantomData<C>);
588
589unsafe impl<'q, C> Fetch<'q> for FetchRead<C>
590where
591    C: 'static,
592{
593    type Ty = u16;
594    type Ref = Ref<'q>;
595    type Ptr = NonNull<C>;
596
597    type Item = &'q C;
598
599    fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
600        borrow_flags.find::<C>()
601    }
602
603    fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
604        archetype.find::<C>()
605    }
606
607    unsafe fn borrow(borrow_flags: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref {
608        borrow_flags.borrow::<C>(ty)
609    }
610
611    unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
612        NonNull::new_unchecked(archetype.base_pointer::<C>(ty))
613    }
614
615    fn dangling() -> Self::Ptr {
616        NonNull::dangling()
617    }
618
619    unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item {
620        &*ptr.as_ptr().add(idx as usize)
621    }
622}
623
624unsafe impl<C> FetchShared for FetchRead<C> {}
625
626impl<'a, C> QuerySpec for &'a mut C
627where
628    C: 'static,
629{
630    type Fetch = FetchWrite<C>;
631}
632
633pub struct FetchWrite<C>(PhantomData<C>);
634
635unsafe impl<'q, C> Fetch<'q> for FetchWrite<C>
636where
637    C: 'static,
638{
639    type Ty = u16;
640    type Ref = RefMut<'q>;
641    type Ptr = NonNull<C>;
642
643    type Item = &'q mut C;
644
645    fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
646        assert_ne!(
647            TypeId::of::<C>(),
648            TypeId::of::<Entity>(),
649            "Entity cannot be queried mutably"
650        );
651
652        borrow_flags.find::<C>()
653    }
654
655    fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
656        archetype.find::<C>()
657    }
658
659    unsafe fn borrow(borrow_flags: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref {
660        borrow_flags.borrow_mut::<C>(ty)
661    }
662
663    unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
664        NonNull::new_unchecked(archetype.base_pointer::<C>(ty))
665    }
666
667    fn dangling() -> Self::Ptr {
668        NonNull::dangling()
669    }
670
671    unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item {
672        &mut *ptr.as_ptr().add(idx as usize)
673    }
674}
675
676impl<S> QuerySpec for Option<S>
677where
678    S: QuerySpec,
679{
680    type Fetch = TryFetch<S::Fetch>;
681}
682
683pub struct TryFetch<F>(PhantomData<F>);
684
685unsafe impl<'q, F> Fetch<'q> for TryFetch<F>
686where
687    F: Fetch<'q>,
688{
689    type Ty = Option<F::Ty>;
690    type Ref = Option<F::Ref>;
691    type Ptr = Option<F::Ptr>;
692
693    type Item = Option<F::Item>;
694
695    fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
696        Some(F::find_flags(borrow_flags))
697    }
698
699    fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
700        Some(F::find_comps(archetype))
701    }
702
703    unsafe fn borrow(borrow_flags: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref {
704        ty.map(|ty| F::borrow(borrow_flags, ty))
705    }
706
707    unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
708        ty.map(|ty| F::base_pointer(archetype, ty))
709    }
710
711    fn dangling() -> Self::Ptr {
712        None
713    }
714
715    unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item {
716        ptr.map(|ptr| F::deref(ptr, idx))
717    }
718}
719
720unsafe impl<F> FetchShared for TryFetch<F> where F: FetchShared {}
721
722/// A query specification to iterate over entities with a certain component,
723/// but without borrowing that component.
724///
725/// # Examples
726///
727/// A query for components of type `u32` and `bool`,
728/// only matching entities with a component of type `f32`.
729///
730/// ```
731/// # use rs_ecs::*;
732/// let query = Query::<With<(&u32, &bool), f32>>::new();
733/// ```
734pub struct With<S, C>(PhantomData<(S, C)>);
735
736impl<S, C> QuerySpec for With<S, C>
737where
738    S: QuerySpec,
739    C: 'static,
740{
741    type Fetch = FetchWith<S::Fetch, C>;
742}
743
744pub struct FetchWith<F, C>(PhantomData<(F, C)>);
745
746unsafe impl<'q, F, C> Fetch<'q> for FetchWith<F, C>
747where
748    F: Fetch<'q>,
749    C: 'static,
750{
751    type Ty = F::Ty;
752    type Ref = F::Ref;
753    type Ptr = F::Ptr;
754
755    type Item = F::Item;
756
757    fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
758        F::find_flags(borrow_flags)
759    }
760
761    fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
762        match archetype.find::<C>() {
763            Some(_) => F::find_comps(archetype),
764            None => None,
765        }
766    }
767
768    unsafe fn borrow(borrow_flags: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref {
769        F::borrow(borrow_flags, ty)
770    }
771
772    unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
773        F::base_pointer(archetype, ty)
774    }
775
776    fn dangling() -> Self::Ptr {
777        F::dangling()
778    }
779
780    unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item {
781        F::deref(ptr, idx)
782    }
783}
784
785unsafe impl<F, C> FetchShared for FetchWith<F, C> where F: FetchShared {}
786
787/// A query specification to iterate over entities without a certain component.
788///
789/// # Examples
790///
791/// A query for components of type `u32` and `bool`,
792/// only matching entities without a component of type `f32`.
793///
794/// ```
795/// # use rs_ecs::*;
796/// let query = Query::<Without<(&u32, &bool), f32>>::new();
797/// ```
798pub struct Without<S, C>(PhantomData<(S, C)>);
799
800impl<S, C> QuerySpec for Without<S, C>
801where
802    S: QuerySpec,
803    C: 'static,
804{
805    type Fetch = FetchWithout<S::Fetch, C>;
806}
807
808pub struct FetchWithout<F, C>(PhantomData<(F, C)>);
809
810unsafe impl<'q, F, C> Fetch<'q> for FetchWithout<F, C>
811where
812    F: Fetch<'q>,
813    C: 'static,
814{
815    type Ty = F::Ty;
816    type Ref = F::Ref;
817    type Ptr = F::Ptr;
818
819    type Item = F::Item;
820
821    fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
822        F::find_flags(borrow_flags)
823    }
824
825    fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
826        match archetype.find::<C>() {
827            None => F::find_comps(archetype),
828            Some(_) => None,
829        }
830    }
831
832    unsafe fn borrow(borrow_flags: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref {
833        F::borrow(borrow_flags, ty)
834    }
835
836    unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
837        F::base_pointer(archetype, ty)
838    }
839
840    fn dangling() -> Self::Ptr {
841        F::dangling()
842    }
843
844    unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item {
845        F::deref(ptr, idx)
846    }
847}
848
849unsafe impl<F, C> FetchShared for FetchWithout<F, C> where F: FetchShared {}
850
851/// A query specification to indicate which entities match the inner query,
852/// but without borrowing any components.
853///
854/// # Examples
855///
856/// ```
857/// # use rs_ecs::*;
858/// let mut world = World::new();
859///
860/// let entity1 = world.alloc();
861/// world.insert(entity1, (42_i32, 1.0_f32));
862///
863/// let entity2 = world.alloc();
864/// world.insert(entity2, (23_i32,));
865///
866/// let mut query = Query::<(&i32, Matches<&f32>)>::new();
867/// let mut query = query.borrow(&world);
868/// let mut query = query.map();
869///
870/// let (i1, f1) = query.get(entity1).unwrap();
871/// assert_eq!(*i1, 42);
872/// assert!(f1);
873///
874/// let (i2, f2) = query.get(entity2).unwrap();
875/// assert_eq!(*i2, 23);
876/// assert!(!f2);
877/// ```
878pub struct Matches<S>(PhantomData<S>);
879
880impl<S> QuerySpec for Matches<S>
881where
882    S: QuerySpec,
883{
884    type Fetch = FetchMatches<S::Fetch>;
885}
886
887pub struct FetchMatches<F>(PhantomData<F>);
888
889unsafe impl<'q, F> Fetch<'q> for FetchMatches<F>
890where
891    F: Fetch<'q>,
892{
893    type Ty = bool;
894    type Ref = ();
895    type Ptr = bool;
896
897    type Item = bool;
898
899    fn find_flags(_borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
900        Some(false)
901    }
902
903    fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
904        Some(F::find_comps(archetype).is_some())
905    }
906
907    unsafe fn borrow(_borrow_flags: &'q BorrowFlags, _ty: Self::Ty) -> Self::Ref {}
908
909    unsafe fn base_pointer(_archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
910        ty
911    }
912
913    fn dangling() -> Self::Ptr {
914        false
915    }
916
917    unsafe fn deref(ptr: Self::Ptr, _idx: u32) -> Self::Item {
918        ptr
919    }
920}
921
922unsafe impl<F> FetchShared for FetchMatches<F> {}
923
924macro_rules! impl_fetch_for_tuples {
925    () => {
926        impl_fetch_for_tuples!(@impl);
927    };
928
929    ($head:ident $(,$tail:ident)*) => {
930        impl_fetch_for_tuples!($($tail),*);
931        impl_fetch_for_tuples!(@rev $head $(,$tail)*;);
932    };
933
934    (@rev ; $($rev:ident),*) => {
935        impl_fetch_for_tuples!(@impl $($rev),*);
936    };
937
938    (@rev $head:ident $(,$tail:ident)*; $($rev:ident),*) => {
939        impl_fetch_for_tuples!(@rev $($tail),*; $head $(,$rev)*);
940    };
941
942    (@impl $($types:ident),*) => {
943        impl<$($types),*> QuerySpec for ($($types,)*)
944        where
945            $($types: QuerySpec,)*
946        {
947            type Fetch = ($($types::Fetch,)*);
948        }
949
950        #[allow(unused_variables, clippy::unused_unit)]
951        unsafe impl<'q, $($types),*> Fetch<'q> for ($($types,)*)
952        where
953            $($types: Fetch<'q>,)*
954        {
955            type Ty = ($($types::Ty,)*);
956            type Ref = ($($types::Ref,)*);
957            type Ptr = ($($types::Ptr,)*);
958
959            type Item = ($($types::Item,)*);
960
961            #[allow(non_snake_case)]
962            fn find_flags(borrow_flags: &BorrowFlags) -> Option<Self::Ty> {
963                $(let $types = $types::find_flags(borrow_flags)?;)*
964
965                Some(($($types,)*))
966            }
967
968            #[allow(non_snake_case)]
969            fn find_comps(archetype: &Archetype) -> Option<Self::Ty> {
970                $(let $types = $types::find_comps(archetype)?;)*
971
972                Some(($($types,)*))
973            }
974
975            #[allow(non_snake_case)]
976            unsafe fn borrow(borrow_flags: &'q BorrowFlags, ty: Self::Ty) -> Self::Ref {
977                let ($($types,)*) = ty;
978
979                ($($types::borrow(borrow_flags, $types),)*)
980            }
981
982            #[allow(non_snake_case)]
983            unsafe fn base_pointer(archetype: &'q Archetype, ty: Self::Ty) -> Self::Ptr {
984                let ($($types,)*) = ty;
985
986                ($($types::base_pointer(archetype, $types),)*)
987            }
988
989            fn dangling() -> Self::Ptr {
990                ($($types::dangling(),)*)
991            }
992
993            #[allow(non_snake_case)]
994            unsafe fn deref(ptr: Self::Ptr, idx: u32) -> Self::Item {
995                let ($($types,)*) = ptr;
996
997                ($($types::deref($types, idx),)*)
998            }
999        }
1000
1001        unsafe impl<$($types),*> FetchShared for ($($types,)*)
1002        where
1003            $($types: FetchShared,)*
1004        {
1005        }
1006    };
1007}
1008
1009impl_fetch_for_tuples!(J, I, H, G, F, E, D, C, B, A);
1010
1011#[cfg(test)]
1012mod tests {
1013    use super::*;
1014
1015    use std::convert::TryInto;
1016    use std::mem::{forget, size_of};
1017    use std::panic::{catch_unwind, AssertUnwindSafe};
1018
1019    fn spawn_three(world: &mut World) {
1020        let ent = world.alloc();
1021        world.insert(ent, (23_i32, 42_u64));
1022
1023        let ent = world.alloc();
1024        world.insert(ent, (1_i32, 2_i8, 3_u16));
1025
1026        let ent = world.alloc();
1027        world.insert(ent, (42_i32, 23_u64, true));
1028    }
1029
1030    #[test]
1031    fn queries_can_be_used_once() {
1032        let mut world = World::new();
1033
1034        spawn_three(&mut world);
1035
1036        let mut query = Query::<&i32>::new();
1037
1038        let comps = query.borrow(&world).iter().copied().collect::<Vec<_>>();
1039        assert_eq!(&comps, &[23, 1, 42]);
1040    }
1041
1042    #[test]
1043    fn queries_can_be_reused() {
1044        let mut world = World::new();
1045
1046        spawn_three(&mut world);
1047
1048        let mut query = Query::<&i32>::new();
1049
1050        let comps1 = query.borrow(&world).iter().copied().collect::<Vec<_>>();
1051        assert_eq!(&comps1, &[23, 1, 42]);
1052
1053        let comps2 = query.borrow(&world).iter().copied().collect::<Vec<_>>();
1054        assert_eq!(&comps2, &comps1);
1055    }
1056
1057    #[test]
1058    fn queries_can_be_reused_after_adding_an_archetype() {
1059        let mut world = World::new();
1060
1061        spawn_three(&mut world);
1062
1063        let mut query = Query::<&i32>::new();
1064
1065        let comps1 = query.borrow(&world).iter().copied().collect::<Vec<_>>();
1066        assert_eq!(&comps1, &[23, 1, 42]);
1067
1068        let ent = world.alloc();
1069        world.insert(ent, (0_i64,));
1070
1071        let comps2 = query.borrow(&world).iter().copied().collect::<Vec<_>>();
1072        assert_eq!(&comps2, &comps1);
1073    }
1074
1075    #[test]
1076    fn queries_can_modify_components() {
1077        let mut world = World::new();
1078
1079        spawn_three(&mut world);
1080
1081        {
1082            let mut query = Query::<&mut i32>::new();
1083
1084            for comp in query.borrow(&world).iter() {
1085                *comp *= -1;
1086            }
1087        }
1088
1089        let mut query = Query::<&i32>::new();
1090        let comps = query.borrow(&world).iter().copied().collect::<Vec<_>>();
1091
1092        assert_eq!(&comps, &[-23, -1, -42]);
1093    }
1094
1095    #[test]
1096    fn forgotten_query_ref_does_not_use_after_free() {
1097        let mut world = World::new();
1098
1099        spawn_three(&mut world);
1100
1101        let mut query = Query::<&mut i32>::new();
1102        forget(query.borrow(&world));
1103
1104        drop(world);
1105        drop(query);
1106    }
1107
1108    #[test]
1109    #[should_panic]
1110    fn forgotten_query_ref_leaks_borrows() {
1111        let mut world = World::new();
1112
1113        spawn_three(&mut world);
1114
1115        let mut query = Query::<&mut i32>::new();
1116        forget(query.borrow(&world));
1117
1118        let _ = query.borrow(&world);
1119    }
1120
1121    #[test]
1122    fn conflicting_borrow_leaves_world_in_consistent_state() {
1123        let mut world = World::new();
1124
1125        spawn_three(&mut world);
1126
1127        let res = catch_unwind(AssertUnwindSafe(|| {
1128            let mut query = Query::<(&i32, Option<(&mut i32, &bool)>)>::new();
1129            let _ = query.borrow(&world);
1130        }));
1131
1132        assert!(res.is_err());
1133
1134        let mut query = Query::<&mut i32>::new();
1135        let _ = query.borrow(&world);
1136    }
1137
1138    #[test]
1139    fn borrows_can_be_reused() {
1140        let mut world = World::new();
1141
1142        spawn_three(&mut world);
1143
1144        let mut query = Query::<&i32>::new();
1145        let mut query = query.borrow(&world);
1146
1147        let cnt1 = query.iter().count();
1148        let cnt2 = query.iter().count();
1149
1150        assert_eq!(cnt1, 3);
1151        assert_eq!(cnt2, cnt1);
1152    }
1153
1154    #[test]
1155    fn entities_can_be_queried_immutably() {
1156        let mut world = World::new();
1157
1158        spawn_three(&mut world);
1159
1160        let mut query = Query::<&Entity>::new();
1161        let cnt = query.borrow(&world).iter().count();
1162
1163        assert_eq!(cnt, 3);
1164    }
1165
1166    #[test]
1167    #[should_panic]
1168    fn entities_cannot_be_queried_mutably() {
1169        let mut world = World::new();
1170
1171        spawn_three(&mut world);
1172
1173        let mut query = Query::<&mut Entity>::new();
1174        let _ = query.borrow(&world).iter();
1175    }
1176
1177    #[test]
1178    fn try_exposes_optional_components() {
1179        let mut world = World::new();
1180
1181        spawn_three(&mut world);
1182
1183        let mut query = Query::<Option<&u64>>::new();
1184        let cnt = query.borrow(&world).iter().count();
1185        let cnt_some = query.borrow(&world).iter().flatten().count();
1186
1187        assert_eq!(cnt, 3);
1188        assert_eq!(cnt_some, 2);
1189    }
1190
1191    #[test]
1192    fn with_checks_for_presence() {
1193        let mut world = World::new();
1194
1195        spawn_three(&mut world);
1196
1197        let mut query = Query::<With<&i32, bool>>::new();
1198        let sum = query.borrow(&world).iter().sum::<i32>();
1199
1200        assert_eq!(sum, 42);
1201    }
1202
1203    #[test]
1204    fn without_checks_for_absence() {
1205        let mut world = World::new();
1206
1207        spawn_three(&mut world);
1208
1209        let mut query = Query::<Without<&i32, bool>>::new();
1210        let sum = query.borrow(&world).iter().sum::<i32>();
1211
1212        assert_eq!(sum, 23 + 1);
1213    }
1214
1215    #[test]
1216    fn map_enables_access_by_entity_id() {
1217        let mut world = World::new();
1218
1219        spawn_three(&mut world);
1220
1221        let entities = Query::<&Entity>::new()
1222            .borrow(&world)
1223            .iter()
1224            .copied()
1225            .collect::<Vec<_>>();
1226
1227        let mut query = Query::<&i32>::new();
1228        let mut query = query.borrow(&world);
1229        let mut query = query.map();
1230
1231        for ent in entities {
1232            query.get_mut(ent).unwrap();
1233        }
1234    }
1235
1236    #[test]
1237    fn map_enables_concurrent_access_to_multiple_entities() {
1238        let mut world = World::new();
1239
1240        spawn_three(&mut world);
1241
1242        let entities = Query::<&Entity>::new()
1243            .borrow(&world)
1244            .iter()
1245            .copied()
1246            .collect::<Vec<_>>();
1247
1248        let mut query = Query::<&mut i32>::new();
1249        let mut query = query.borrow(&world);
1250        let mut query = query.map();
1251
1252        let entities: [Entity; 3] = entities.try_into().unwrap();
1253
1254        let mut values = query.get_many_mut(entities);
1255
1256        for (index, value) in values.iter_mut().enumerate() {
1257            **value.as_mut().unwrap() = index as i32;
1258        }
1259    }
1260
1261    #[test]
1262    #[should_panic]
1263    fn map_prohibits_aliasing_via_concurrent_access() {
1264        let mut world = World::new();
1265
1266        spawn_three(&mut world);
1267
1268        let entities = Query::<&Entity>::new()
1269            .borrow(&world)
1270            .iter()
1271            .copied()
1272            .collect::<Vec<_>>();
1273
1274        let mut query = Query::<&mut i32>::new();
1275        let mut query = query.borrow(&world);
1276        let mut query = query.map();
1277
1278        let entities = [entities[0], entities[1], entities[1]];
1279
1280        query.get_many_mut(entities);
1281    }
1282
1283    #[test]
1284    fn empty_queries_yields_unit_for_all_entities() {
1285        let mut world = World::new();
1286
1287        spawn_three(&mut world);
1288
1289        let res = Query::<()>::new().borrow(&world).iter().collect::<Vec<_>>();
1290
1291        assert_eq!(res, vec![(), (), ()]);
1292    }
1293
1294    #[test]
1295    fn fetch_ptr_has_niche() {
1296        assert_eq!(
1297            size_of::<<<(&i32, &mut f64) as QuerySpec>::Fetch as Fetch>::Ptr>(),
1298            size_of::<Option<<<(&i32, &mut f64) as QuerySpec>::Fetch as Fetch>::Ptr>>()
1299        );
1300    }
1301}
rs_ecs/query.rs

rs_ecs/
query.rs
