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

rs_ecs/
query.rs
