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

rs_ecs/
query.rs
