hv_ecs/
query.rs

1// Copyright 2019 Google LLC
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8use core::any::TypeId;
9use core::marker::PhantomData;
10use core::mem;
11use core::ptr::NonNull;
12use core::slice::Iter as SliceIter;
13
14use crate::alloc::boxed::Box;
15use crate::archetype::Archetype;
16use crate::entities::EntityMeta;
17use crate::{Component, Entity, World};
18
19/// A collection of component types to fetch from a [`World`](crate::World)
20pub trait Query {
21    #[doc(hidden)]
22    type Fetch: for<'a> Fetch<'a>;
23}
24
25/// Type of values yielded by a query
26///
27/// Once rust offers generic associated types, this will be moved into [`Query`].
28pub type QueryItem<'a, Q> = <<Q as Query>::Fetch as Fetch<'a>>::Item;
29
30/// Streaming iterators over contiguous homogeneous ranges of components
31pub unsafe trait Fetch<'a>: Sized {
32    /// Type of value to be fetched
33    type Item;
34
35    /// The type of the data which can be cached to speed up retrieving
36    /// the relevant type states from a matching [`Archetype`]
37    type State: Copy;
38
39    /// A value on which `get` may never be called
40    fn dangling() -> Self;
41
42    /// How this query will access `archetype`, if at all
43    fn access(archetype: &Archetype) -> Option<Access>;
44
45    /// Acquire dynamic borrows from `archetype`
46    fn borrow(archetype: &Archetype, state: Self::State);
47    /// Look up state for `archetype` if it should be traversed
48    fn prepare(archetype: &Archetype) -> Option<Self::State>;
49    /// Construct a `Fetch` for `archetype` based on the associated state
50    fn execute(archetype: &'a Archetype, state: Self::State) -> Self;
51    /// Release dynamic borrows acquired by `borrow`
52    fn release(archetype: &Archetype, state: Self::State);
53
54    /// Invoke `f` for every component type that may be borrowed and whether the borrow is unique
55    fn for_each_borrow(f: impl FnMut(TypeId, bool));
56
57    /// Access the `n`th item in this archetype without bounds checking
58    ///
59    /// # Safety
60    /// - Must only be called after `borrow`
61    /// - `release` must not be called while `'a` is still live
62    /// - Bounds-checking must be performed externally
63    /// - Any resulting borrows must be legal (e.g. no &mut to something another iterator might access)
64    unsafe fn get(&self, n: usize) -> Self::Item;
65}
66
67/// Type of access a [`Query`] may have to an [`Archetype`]
68#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
69pub enum Access {
70    /// Read entity IDs only, no components
71    Iterate,
72    /// Read components
73    Read,
74    /// Read and write components
75    Write,
76}
77
78impl<'a, T: Component> Query for &'a T {
79    type Fetch = FetchRead<T>;
80}
81
82#[doc(hidden)]
83pub struct FetchRead<T>(NonNull<T>);
84
85unsafe impl<'a, T: Component> Fetch<'a> for FetchRead<T> {
86    type Item = &'a T;
87
88    type State = usize;
89
90    fn dangling() -> Self {
91        Self(NonNull::dangling())
92    }
93
94    fn access(archetype: &Archetype) -> Option<Access> {
95        if archetype.has::<T>() {
96            Some(Access::Read)
97        } else {
98            None
99        }
100    }
101
102    fn borrow(archetype: &Archetype, state: Self::State) {
103        archetype.borrow::<T>(state);
104    }
105    fn prepare(archetype: &Archetype) -> Option<Self::State> {
106        archetype.get_state::<T>()
107    }
108    fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
109        Self(archetype.get_base(state))
110    }
111    fn release(archetype: &Archetype, state: Self::State) {
112        archetype.release::<T>(state);
113    }
114
115    fn for_each_borrow(mut f: impl FnMut(TypeId, bool)) {
116        f(TypeId::of::<T>(), false);
117    }
118
119    unsafe fn get(&self, n: usize) -> Self::Item {
120        &*self.0.as_ptr().add(n)
121    }
122}
123
124impl<'a, T: Component> Query for &'a mut T {
125    type Fetch = FetchWrite<T>;
126}
127
128#[doc(hidden)]
129pub struct FetchWrite<T>(NonNull<T>);
130
131unsafe impl<'a, T: Component> Fetch<'a> for FetchWrite<T> {
132    type Item = &'a mut T;
133
134    type State = usize;
135
136    fn dangling() -> Self {
137        Self(NonNull::dangling())
138    }
139
140    fn access(archetype: &Archetype) -> Option<Access> {
141        if archetype.has::<T>() {
142            Some(Access::Write)
143        } else {
144            None
145        }
146    }
147
148    fn borrow(archetype: &Archetype, state: Self::State) {
149        archetype.borrow_mut::<T>(state);
150    }
151    #[allow(clippy::needless_question_mark)]
152    fn prepare(archetype: &Archetype) -> Option<Self::State> {
153        Some(archetype.get_state::<T>()?)
154    }
155    fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
156        Self(archetype.get_base::<T>(state))
157    }
158    fn release(archetype: &Archetype, state: Self::State) {
159        archetype.release_mut::<T>(state);
160    }
161
162    fn for_each_borrow(mut f: impl FnMut(TypeId, bool)) {
163        f(TypeId::of::<T>(), true);
164    }
165
166    unsafe fn get(&self, n: usize) -> Self::Item {
167        &mut *self.0.as_ptr().add(n)
168    }
169}
170
171impl<T: Query> Query for Option<T> {
172    type Fetch = TryFetch<T::Fetch>;
173}
174
175#[doc(hidden)]
176pub struct TryFetch<T>(Option<T>);
177
178unsafe impl<'a, T: Fetch<'a>> Fetch<'a> for TryFetch<T> {
179    type Item = Option<T::Item>;
180
181    type State = Option<T::State>;
182
183    fn dangling() -> Self {
184        Self(None)
185    }
186
187    fn access(archetype: &Archetype) -> Option<Access> {
188        Some(T::access(archetype).unwrap_or(Access::Iterate))
189    }
190
191    fn borrow(archetype: &Archetype, state: Self::State) {
192        if let Some(state) = state {
193            T::borrow(archetype, state);
194        }
195    }
196    fn prepare(archetype: &Archetype) -> Option<Self::State> {
197        Some(T::prepare(archetype))
198    }
199    fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
200        Self(state.map(|state| T::execute(archetype, state)))
201    }
202    fn release(archetype: &Archetype, state: Self::State) {
203        if let Some(state) = state {
204            T::release(archetype, state);
205        }
206    }
207
208    fn for_each_borrow(f: impl FnMut(TypeId, bool)) {
209        T::for_each_borrow(f);
210    }
211
212    unsafe fn get(&self, n: usize) -> Option<T::Item> {
213        Some(self.0.as_ref()?.get(n))
214    }
215}
216
217/// Holds an `L`, or an `R`, or both
218#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
219pub enum Or<L, R> {
220    /// Just an `L`
221    Left(L),
222    /// Just an `R`
223    Right(R),
224    /// Both an `L` and an `R`
225    Both(L, R),
226}
227
228impl<L, R> Or<L, R> {
229    /// Construct an `Or<L, R>` if at least one argument is `Some`
230    pub fn new(l: Option<L>, r: Option<R>) -> Option<Self> {
231        match (l, r) {
232            (None, None) => None,
233            (Some(l), None) => Some(Or::Left(l)),
234            (None, Some(r)) => Some(Or::Right(r)),
235            (Some(l), Some(r)) => Some(Or::Both(l, r)),
236        }
237    }
238
239    /// Destructure into two `Option`s, where either or both are `Some`
240    pub fn split(self) -> (Option<L>, Option<R>) {
241        match self {
242            Or::Left(l) => (Some(l), None),
243            Or::Right(r) => (None, Some(r)),
244            Or::Both(l, r) => (Some(l), Some(r)),
245        }
246    }
247
248    /// Extract `L` regardless of whether `R` is present
249    pub fn left(self) -> Option<L> {
250        match self {
251            Or::Left(l) => Some(l),
252            Or::Both(l, _) => Some(l),
253            _ => None,
254        }
255    }
256
257    /// Extract `R` regardless of whether `L` is present
258    pub fn right(self) -> Option<R> {
259        match self {
260            Or::Right(r) => Some(r),
261            Or::Both(_, r) => Some(r),
262            _ => None,
263        }
264    }
265
266    /// Transform `L` with `f` and `R` with `g`
267    pub fn map<L1, R1, F, G>(self, f: F, g: G) -> Or<L1, R1>
268    where
269        F: FnOnce(L) -> L1,
270        G: FnOnce(R) -> R1,
271    {
272        match self {
273            Or::Left(l) => Or::Left(f(l)),
274            Or::Right(r) => Or::Right(g(r)),
275            Or::Both(l, r) => Or::Both(f(l), g(r)),
276        }
277    }
278
279    /// Convert from `&Or<L, R>` to `Or<&L, &R>`
280    pub fn as_ref(&self) -> Or<&L, &R> {
281        match *self {
282            Or::Left(ref l) => Or::Left(l),
283            Or::Right(ref r) => Or::Right(r),
284            Or::Both(ref l, ref r) => Or::Both(l, r),
285        }
286    }
287
288    /// Convert from `&mut Or<L, R>` to `Or<&mut L, &mut R>`
289    pub fn as_mut(&mut self) -> Or<&mut L, &mut R> {
290        match *self {
291            Or::Left(ref mut l) => Or::Left(l),
292            Or::Right(ref mut r) => Or::Right(r),
293            Or::Both(ref mut l, ref mut r) => Or::Both(l, r),
294        }
295    }
296}
297
298impl<L, R> Or<&'_ L, &'_ R>
299where
300    L: Clone,
301    R: Clone,
302{
303    /// Maps an `Or<&L, &R>` to an `Or<L, R>` by cloning its contents
304    pub fn cloned(self) -> Or<L, R> {
305        self.map(Clone::clone, Clone::clone)
306    }
307}
308
309impl<L: Query, R: Query> Query for Or<L, R> {
310    type Fetch = FetchOr<L::Fetch, R::Fetch>;
311}
312
313#[doc(hidden)]
314pub struct FetchOr<L, R>(Or<L, R>);
315
316unsafe impl<'a, L: Fetch<'a>, R: Fetch<'a>> Fetch<'a> for FetchOr<L, R> {
317    type Item = Or<L::Item, R::Item>;
318
319    type State = Or<L::State, R::State>;
320
321    fn dangling() -> Self {
322        Self(Or::Left(L::dangling()))
323    }
324
325    fn access(archetype: &Archetype) -> Option<Access> {
326        L::access(archetype).max(R::access(archetype))
327    }
328
329    fn borrow(archetype: &Archetype, state: Self::State) {
330        state.map(|l| L::borrow(archetype, l), |r| R::borrow(archetype, r));
331    }
332
333    fn prepare(archetype: &Archetype) -> Option<Self::State> {
334        Or::new(L::prepare(archetype), R::prepare(archetype))
335    }
336
337    fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
338        Self(state.map(|l| L::execute(archetype, l), |r| R::execute(archetype, r)))
339    }
340
341    fn release(archetype: &Archetype, state: Self::State) {
342        state.map(|l| L::release(archetype, l), |r| R::release(archetype, r));
343    }
344
345    fn for_each_borrow(mut f: impl FnMut(TypeId, bool)) {
346        L::for_each_borrow(&mut f);
347        R::for_each_borrow(&mut f);
348    }
349
350    unsafe fn get(&self, n: usize) -> Self::Item {
351        self.0.as_ref().map(|l| l.get(n), |r| r.get(n))
352    }
353}
354
355/// Query transformer skipping entities that have a `T` component
356///
357/// See also `QueryBorrow::without`.
358///
359/// # Example
360/// ```
361/// # use hecs::*;
362/// let mut world = World::new();
363/// let a = world.spawn((123, true, "abc"));
364/// let b = world.spawn((456, false));
365/// let c = world.spawn((42, "def"));
366/// let entities = world.query::<Without<bool, &i32>>()
367///     .iter()
368///     .map(|(e, &i)| (e, i))
369///     .collect::<Vec<_>>();
370/// assert_eq!(entities, &[(c, 42)]);
371/// ```
372pub struct Without<T, Q>(PhantomData<(Q, fn(T))>);
373
374impl<T: Component, Q: Query> Query for Without<T, Q> {
375    type Fetch = FetchWithout<T, Q::Fetch>;
376}
377
378#[doc(hidden)]
379pub struct FetchWithout<T, F>(F, PhantomData<fn(T)>);
380
381unsafe impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWithout<T, F> {
382    type Item = F::Item;
383
384    type State = F::State;
385
386    fn dangling() -> Self {
387        Self(F::dangling(), PhantomData)
388    }
389
390    fn access(archetype: &Archetype) -> Option<Access> {
391        if archetype.has::<T>() {
392            None
393        } else {
394            F::access(archetype)
395        }
396    }
397
398    fn borrow(archetype: &Archetype, state: Self::State) {
399        F::borrow(archetype, state)
400    }
401    fn prepare(archetype: &Archetype) -> Option<Self::State> {
402        if archetype.has::<T>() {
403            return None;
404        }
405        F::prepare(archetype)
406    }
407    fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
408        Self(F::execute(archetype, state), PhantomData)
409    }
410    fn release(archetype: &Archetype, state: Self::State) {
411        F::release(archetype, state)
412    }
413
414    fn for_each_borrow(f: impl FnMut(TypeId, bool)) {
415        F::for_each_borrow(f);
416    }
417
418    unsafe fn get(&self, n: usize) -> F::Item {
419        self.0.get(n)
420    }
421}
422
423/// Query transformer skipping entities that do not have a `T` component
424///
425/// See also `QueryBorrow::with`.
426///
427/// # Example
428/// ```
429/// # use hecs::*;
430/// let mut world = World::new();
431/// let a = world.spawn((123, true, "abc"));
432/// let b = world.spawn((456, false));
433/// let c = world.spawn((42, "def"));
434/// let entities = world.query::<With<bool, &i32>>()
435///     .iter()
436///     .map(|(e, &i)| (e, i))
437///     .collect::<Vec<_>>();
438/// assert_eq!(entities.len(), 2);
439/// assert!(entities.contains(&(a, 123)));
440/// assert!(entities.contains(&(b, 456)));
441/// ```
442pub struct With<T, Q>(PhantomData<(Q, fn(T))>);
443
444impl<T: Component, Q: Query> Query for With<T, Q> {
445    type Fetch = FetchWith<T, Q::Fetch>;
446}
447
448#[doc(hidden)]
449pub struct FetchWith<T, F>(F, PhantomData<fn(T)>);
450
451unsafe impl<'a, T: Component, F: Fetch<'a>> Fetch<'a> for FetchWith<T, F> {
452    type Item = F::Item;
453
454    type State = F::State;
455
456    fn dangling() -> Self {
457        Self(F::dangling(), PhantomData)
458    }
459
460    fn access(archetype: &Archetype) -> Option<Access> {
461        if archetype.has::<T>() {
462            F::access(archetype)
463        } else {
464            None
465        }
466    }
467
468    fn borrow(archetype: &Archetype, state: Self::State) {
469        F::borrow(archetype, state)
470    }
471    fn prepare(archetype: &Archetype) -> Option<Self::State> {
472        if !archetype.has::<T>() {
473            return None;
474        }
475        F::prepare(archetype)
476    }
477    fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
478        Self(F::execute(archetype, state), PhantomData)
479    }
480    fn release(archetype: &Archetype, state: Self::State) {
481        F::release(archetype, state)
482    }
483
484    fn for_each_borrow(f: impl FnMut(TypeId, bool)) {
485        F::for_each_borrow(f);
486    }
487
488    unsafe fn get(&self, n: usize) -> F::Item {
489        self.0.get(n)
490    }
491}
492
493/// A query that yields `true` iff an entity would satisfy the query `Q`
494///
495/// Does not borrow any components, making it faster and more concurrency-friendly than `Option<Q>`.
496///
497/// # Example
498/// ```
499/// # use hecs::*;
500/// let mut world = World::new();
501/// let a = world.spawn((123, true, "abc"));
502/// let b = world.spawn((456, false));
503/// let c = world.spawn((42, "def"));
504/// let entities = world.query::<Satisfies<&bool>>()
505///     .iter()
506///     .map(|(e, x)| (e, x))
507///     .collect::<Vec<_>>();
508/// assert_eq!(entities.len(), 3);
509/// assert!(entities.contains(&(a, true)));
510/// assert!(entities.contains(&(b, true)));
511/// assert!(entities.contains(&(c, false)));
512/// ```
513pub struct Satisfies<Q>(PhantomData<Q>);
514
515impl<Q: Query> Query for Satisfies<Q> {
516    type Fetch = FetchSatisfies<Q::Fetch>;
517}
518
519#[doc(hidden)]
520pub struct FetchSatisfies<F>(bool, PhantomData<F>);
521
522unsafe impl<'a, F: Fetch<'a>> Fetch<'a> for FetchSatisfies<F> {
523    type Item = bool;
524
525    type State = bool;
526
527    fn dangling() -> Self {
528        Self(false, PhantomData)
529    }
530
531    fn access(archetype: &Archetype) -> Option<Access> {
532        F::access(archetype).map(|_| Access::Iterate)
533    }
534
535    fn borrow(_archetype: &Archetype, _state: Self::State) {}
536    fn prepare(archetype: &Archetype) -> Option<Self::State> {
537        Some(F::prepare(archetype).is_some())
538    }
539    fn execute(_archetype: &'a Archetype, state: Self::State) -> Self {
540        Self(state, PhantomData)
541    }
542    fn release(_archetype: &Archetype, _state: Self::State) {}
543
544    fn for_each_borrow(_: impl FnMut(TypeId, bool)) {}
545
546    unsafe fn get(&self, _: usize) -> bool {
547        self.0
548    }
549}
550
551/// A borrow of a [`World`](crate::World) sufficient to execute the query `Q`
552///
553/// Note that borrows are not released until this object is dropped.
554pub struct QueryBorrow<'w, Q: Query> {
555    meta: &'w [EntityMeta],
556    archetypes: &'w [Archetype],
557    borrowed: bool,
558    _marker: PhantomData<Q>,
559}
560
561impl<'w, Q: Query> QueryBorrow<'w, Q> {
562    pub(crate) fn new(meta: &'w [EntityMeta], archetypes: &'w [Archetype]) -> Self {
563        Self {
564            meta,
565            archetypes,
566            borrowed: false,
567            _marker: PhantomData,
568        }
569    }
570
571    /// Execute the query
572    ///
573    /// Must be called only once per query.
574    // The lifetime narrowing here is required for soundness.
575    pub fn iter(&mut self) -> QueryIter<'_, Q> {
576        self.borrow();
577        unsafe { QueryIter::new(self.meta, self.archetypes.iter()) }
578    }
579
580    /// Like `iter`, but returns child iterators of at most `batch_size` elements
581    ///
582    /// Useful for distributing work over a threadpool.
583    // The lifetime narrowing here is required for soundness.
584    pub fn iter_batched(&mut self, batch_size: u32) -> BatchedIter<'_, Q> {
585        self.borrow();
586        unsafe { BatchedIter::new(self.meta, self.archetypes.iter(), batch_size) }
587    }
588
589    fn borrow(&mut self) {
590        if self.borrowed {
591            return;
592        }
593        for x in self.archetypes {
594            // TODO: Release prior borrows on failure?
595            if let Some(state) = Q::Fetch::prepare(x) {
596                Q::Fetch::borrow(x, state);
597            }
598        }
599        self.borrowed = true;
600    }
601
602    /// Transform the query into one that requires a certain component without borrowing it
603    ///
604    /// This can be useful when the component needs to be borrowed elsewhere and it isn't necessary
605    /// for the iterator to expose its data directly.
606    ///
607    /// Equivalent to using a query type wrapped in `With`.
608    ///
609    /// # Example
610    /// ```
611    /// # use hecs::*;
612    /// let mut world = World::new();
613    /// let a = world.spawn((123, true, "abc"));
614    /// let b = world.spawn((456, false));
615    /// let c = world.spawn((42, "def"));
616    /// let entities = world.query::<&i32>()
617    ///     .with::<bool>()
618    ///     .iter()
619    ///     .map(|(e, &i)| (e, i)) // Copy out of the world
620    ///     .collect::<Vec<_>>();
621    /// assert!(entities.contains(&(a, 123)));
622    /// assert!(entities.contains(&(b, 456)));
623    /// ```
624    pub fn with<T: Component>(self) -> QueryBorrow<'w, With<T, Q>> {
625        self.transform()
626    }
627
628    /// Transform the query into one that skips entities having a certain component
629    ///
630    /// Equivalent to using a query type wrapped in `Without`.
631    ///
632    /// # Example
633    /// ```
634    /// # use hecs::*;
635    /// let mut world = World::new();
636    /// let a = world.spawn((123, true, "abc"));
637    /// let b = world.spawn((456, false));
638    /// let c = world.spawn((42, "def"));
639    /// let entities = world.query::<&i32>()
640    ///     .without::<bool>()
641    ///     .iter()
642    ///     .map(|(e, &i)| (e, i)) // Copy out of the world
643    ///     .collect::<Vec<_>>();
644    /// assert_eq!(entities, &[(c, 42)]);
645    /// ```
646    pub fn without<T: Component>(self) -> QueryBorrow<'w, Without<T, Q>> {
647        self.transform()
648    }
649
650    /// Helper to change the type of the query
651    fn transform<R: Query>(mut self) -> QueryBorrow<'w, R> {
652        let x = QueryBorrow {
653            meta: self.meta,
654            archetypes: self.archetypes,
655            borrowed: self.borrowed,
656            _marker: PhantomData,
657        };
658        // Ensure `Drop` won't fire redundantly
659        self.borrowed = false;
660        x
661    }
662}
663
664unsafe impl<'w, Q: Query> Send for QueryBorrow<'w, Q> {}
665unsafe impl<'w, Q: Query> Sync for QueryBorrow<'w, Q> {}
666
667impl<'w, Q: Query> Drop for QueryBorrow<'w, Q> {
668    fn drop(&mut self) {
669        if self.borrowed {
670            for x in self.archetypes {
671                if let Some(state) = Q::Fetch::prepare(x) {
672                    Q::Fetch::release(x, state);
673                }
674            }
675        }
676    }
677}
678
679impl<'q, 'w, Q: Query> IntoIterator for &'q mut QueryBorrow<'w, Q> {
680    type Item = (Entity, QueryItem<'q, Q>);
681    type IntoIter = QueryIter<'q, Q>;
682
683    fn into_iter(self) -> Self::IntoIter {
684        self.iter()
685    }
686}
687
688/// Iterator over the set of entities with the components in `Q`
689pub struct QueryIter<'q, Q: Query> {
690    meta: &'q [EntityMeta],
691    archetypes: SliceIter<'q, Archetype>,
692    iter: ChunkIter<Q>,
693}
694
695impl<'q, Q: Query> QueryIter<'q, Q> {
696    /// # Safety
697    ///
698    /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
699    /// dynamic borrow checks or by representing exclusive access to the `World`.
700    unsafe fn new(meta: &'q [EntityMeta], archetypes: SliceIter<'q, Archetype>) -> Self {
701        Self {
702            meta,
703            archetypes,
704            iter: ChunkIter::empty(),
705        }
706    }
707}
708
709unsafe impl<'q, Q: Query> Send for QueryIter<'q, Q> {}
710unsafe impl<'q, Q: Query> Sync for QueryIter<'q, Q> {}
711
712impl<'q, Q: Query> Iterator for QueryIter<'q, Q> {
713    type Item = (Entity, QueryItem<'q, Q>);
714
715    #[inline(always)]
716    fn next(&mut self) -> Option<Self::Item> {
717        loop {
718            match unsafe { self.iter.next() } {
719                None => {
720                    let archetype = self.archetypes.next()?;
721                    let state = Q::Fetch::prepare(archetype);
722                    let fetch = state.map(|state| Q::Fetch::execute(archetype, state));
723                    self.iter = fetch.map_or(ChunkIter::empty(), |fetch| ChunkIter {
724                        entities: archetype.entities(),
725                        fetch,
726                        position: 0,
727                        len: archetype.len() as usize,
728                    });
729                    continue;
730                }
731                Some((id, components)) => {
732                    return Some((
733                        Entity {
734                            id,
735                            generation: unsafe { self.meta.get_unchecked(id as usize).generation },
736                        },
737                        components,
738                    ));
739                }
740            }
741        }
742    }
743
744    fn size_hint(&self) -> (usize, Option<usize>) {
745        let n = self.len();
746        (n, Some(n))
747    }
748}
749
750impl<'q, Q: Query> ExactSizeIterator for QueryIter<'q, Q> {
751    fn len(&self) -> usize {
752        self.archetypes
753            .clone()
754            .filter(|&x| Q::Fetch::access(x).is_some())
755            .map(|x| x.len() as usize)
756            .sum::<usize>()
757            + self.iter.remaining()
758    }
759}
760
761/// A query builder that's convertible directly into an iterator
762pub struct QueryMut<'q, Q: Query> {
763    iter: QueryIter<'q, Q>,
764}
765
766impl<'q, Q: Query> QueryMut<'q, Q> {
767    pub(crate) fn new(meta: &'q [EntityMeta], archetypes: &'q mut [Archetype]) -> Self {
768        assert_borrow::<Q>();
769
770        Self {
771            iter: unsafe { QueryIter::new(meta, archetypes.iter()) },
772        }
773    }
774
775    /// Transform the query into one that requires a certain component without borrowing it
776    ///
777    /// See `QueryBorrow::with`
778    pub fn with<T: Component>(self) -> QueryMut<'q, With<T, Q>> {
779        self.transform()
780    }
781
782    /// Transform the query into one that skips entities having a certain component
783    ///
784    /// See `QueryBorrow::without`
785    pub fn without<T: Component>(self) -> QueryMut<'q, Without<T, Q>> {
786        self.transform()
787    }
788
789    /// Helper to change the type of the query
790    fn transform<R: Query>(self) -> QueryMut<'q, R> {
791        QueryMut {
792            iter: unsafe { QueryIter::new(self.iter.meta, self.iter.archetypes) },
793        }
794    }
795}
796
797impl<'q, Q: Query> IntoIterator for QueryMut<'q, Q> {
798    type Item = <QueryIter<'q, Q> as Iterator>::Item;
799    type IntoIter = QueryIter<'q, Q>;
800
801    #[inline]
802    fn into_iter(self) -> Self::IntoIter {
803        self.iter
804    }
805}
806
807fn assert_borrow<Q: Query>() {
808    // This looks like an ugly O(n^2) loop, but everything's constant after inlining, so in
809    // practice LLVM optimizes it out entirely.
810    let mut i = 0;
811    Q::Fetch::for_each_borrow(|a, unique| {
812        if unique {
813            let mut j = 0;
814            Q::Fetch::for_each_borrow(|b, _| {
815                if i != j {
816                    core::assert!(a != b, "query violates a unique borrow");
817                }
818                j += 1;
819            })
820        }
821        i += 1;
822    });
823}
824
825struct ChunkIter<Q: Query> {
826    entities: NonNull<u32>,
827    fetch: Q::Fetch,
828    position: usize,
829    len: usize,
830}
831
832impl<Q: Query> ChunkIter<Q> {
833    fn empty() -> Self {
834        Self {
835            entities: NonNull::dangling(),
836            fetch: Q::Fetch::dangling(),
837            position: 0,
838            len: 0,
839        }
840    }
841
842    #[inline]
843    unsafe fn next<'a>(&mut self) -> Option<(u32, <Q::Fetch as Fetch<'a>>::Item)> {
844        if self.position == self.len {
845            return None;
846        }
847        let entity = self.entities.as_ptr().add(self.position);
848        let item = self.fetch.get(self.position);
849        self.position += 1;
850        Some((*entity, item))
851    }
852
853    fn remaining(&self) -> usize {
854        self.len - self.position
855    }
856}
857
858/// Batched version of [`QueryIter`]
859pub struct BatchedIter<'q, Q: Query> {
860    _marker: PhantomData<&'q Q>,
861    meta: &'q [EntityMeta],
862    archetypes: SliceIter<'q, Archetype>,
863    batch_size: u32,
864    batch: u32,
865}
866
867impl<'q, Q: Query> BatchedIter<'q, Q> {
868    /// # Safety
869    ///
870    /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
871    /// dynamic borrow checks or by representing exclusive access to the `World`.
872    unsafe fn new(
873        meta: &'q [EntityMeta],
874        archetypes: SliceIter<'q, Archetype>,
875        batch_size: u32,
876    ) -> Self {
877        Self {
878            _marker: PhantomData,
879            meta,
880            archetypes,
881            batch_size,
882            batch: 0,
883        }
884    }
885}
886
887unsafe impl<'q, Q: Query> Send for BatchedIter<'q, Q> {}
888unsafe impl<'q, Q: Query> Sync for BatchedIter<'q, Q> {}
889
890impl<'q, Q: Query> Iterator for BatchedIter<'q, Q> {
891    type Item = Batch<'q, Q>;
892
893    fn next(&mut self) -> Option<Self::Item> {
894        loop {
895            let mut archetypes = self.archetypes.clone();
896            let archetype = archetypes.next()?;
897            let offset = self.batch_size * self.batch;
898            if offset >= archetype.len() {
899                self.archetypes = archetypes;
900                self.batch = 0;
901                continue;
902            }
903            let state = Q::Fetch::prepare(archetype);
904            let fetch = state.map(|state| Q::Fetch::execute(archetype, state));
905            if let Some(fetch) = fetch {
906                self.batch += 1;
907                return Some(Batch {
908                    meta: self.meta,
909                    state: ChunkIter {
910                        entities: archetype.entities(),
911                        fetch,
912                        len: (offset + self.batch_size.min(archetype.len() - offset)) as usize,
913                        position: offset as usize,
914                    },
915                });
916            } else {
917                self.archetypes = archetypes;
918                debug_assert_eq!(
919                    self.batch, 0,
920                    "query fetch should always reject at the first batch or not at all"
921                );
922                continue;
923            }
924        }
925    }
926}
927
928/// A sequence of entities yielded by `BatchedIter`
929pub struct Batch<'q, Q: Query> {
930    meta: &'q [EntityMeta],
931    state: ChunkIter<Q>,
932}
933
934impl<'q, Q: Query> Iterator for Batch<'q, Q> {
935    type Item = (Entity, QueryItem<'q, Q>);
936
937    fn next(&mut self) -> Option<Self::Item> {
938        let (id, components) = unsafe { self.state.next()? };
939        Some((
940            Entity {
941                id,
942                generation: self.meta[id as usize].generation,
943            },
944            components,
945        ))
946    }
947}
948
949unsafe impl<'q, Q: Query> Send for Batch<'q, Q> {}
950unsafe impl<'q, Q: Query> Sync for Batch<'q, Q> {}
951
952macro_rules! tuple_impl {
953    ($($name: ident),*) => {
954        unsafe impl<'a, $($name: Fetch<'a>),*> Fetch<'a> for ($($name,)*) {
955            type Item = ($($name::Item,)*);
956
957            type State = ($($name::State,)*);
958
959            #[allow(clippy::unused_unit)]
960            fn dangling() -> Self {
961                ($($name::dangling(),)*)
962            }
963
964            #[allow(unused_variables, unused_mut)]
965            fn access(archetype: &Archetype) -> Option<Access> {
966                let mut access = Access::Iterate;
967                $(
968                    access = access.max($name::access(archetype)?);
969                )*
970                Some(access)
971            }
972
973            #[allow(unused_variables, non_snake_case, clippy::unused_unit)]
974            fn borrow(archetype: &Archetype, state: Self::State) {
975                let ($($name,)*) = state;
976                $($name::borrow(archetype, $name);)*
977            }
978            #[allow(unused_variables)]
979            fn prepare(archetype: &Archetype) -> Option<Self::State> {
980                Some(($($name::prepare(archetype)?,)*))
981            }
982            #[allow(unused_variables, non_snake_case, clippy::unused_unit)]
983            fn execute(archetype: &'a Archetype, state: Self::State) -> Self {
984                let ($($name,)*) = state;
985                ($($name::execute(archetype, $name),)*)
986            }
987            #[allow(unused_variables, non_snake_case, clippy::unused_unit)]
988            fn release(archetype: &Archetype, state: Self::State) {
989                let ($($name,)*) = state;
990                $($name::release(archetype, $name);)*
991            }
992
993            #[allow(unused_variables, unused_mut, clippy::unused_unit)]
994            fn for_each_borrow(mut f: impl FnMut(TypeId, bool)) {
995                $($name::for_each_borrow(&mut f);)*
996            }
997
998            #[allow(unused_variables, clippy::unused_unit)]
999            unsafe fn get(&self, n: usize) -> Self::Item {
1000                #[allow(non_snake_case)]
1001                let ($($name,)*) = self;
1002                ($($name.get(n),)*)
1003            }
1004        }
1005
1006        impl<$($name: Query),*> Query for ($($name,)*) {
1007            type Fetch = ($($name::Fetch,)*);
1008        }
1009    };
1010}
1011
1012//smaller_tuples_too!(tuple_impl, B, A);
1013smaller_tuples_too!(tuple_impl, O, N, M, L, K, J, I, H, G, F, E, D, C, B, A);
1014
1015/// A prepared query can be stored independently of the [`World`] to amortize query set-up costs.
1016pub struct PreparedQuery<Q: Query> {
1017    memo: (u64, u64),
1018    state: Box<[(usize, <Q::Fetch as Fetch<'static>>::State)]>,
1019}
1020
1021impl<Q: Query> Default for PreparedQuery<Q> {
1022    fn default() -> Self {
1023        Self::new()
1024    }
1025}
1026
1027impl<Q: Query> PreparedQuery<Q> {
1028    /// Create a prepared query which is not yet attached to any world
1029    pub fn new() -> Self {
1030        Self {
1031            // This memo will not match any world as the first ID will 1.
1032            memo: (0, 0),
1033            state: Default::default(),
1034        }
1035    }
1036
1037    #[cold]
1038    fn prepare(world: &World) -> Self {
1039        let memo = world.memo();
1040
1041        let state = world
1042            .archetypes()
1043            .enumerate()
1044            .filter_map(|(idx, x)| Q::Fetch::prepare(x).map(|state| (idx, state)))
1045            .collect();
1046
1047        Self { memo, state }
1048    }
1049
1050    /// Query `world`, using dynamic borrow checking
1051    ///
1052    /// This will panic if it would violate an existing unique reference
1053    /// or construct an invalid unique reference.
1054    pub fn query<'q>(&'q mut self, world: &'q World) -> PreparedQueryBorrow<'q, Q> {
1055        if self.memo != world.memo() {
1056            *self = Self::prepare(world);
1057        }
1058
1059        let meta = world.entities_meta();
1060        let archetypes = world.archetypes_inner();
1061
1062        PreparedQueryBorrow::new(meta, archetypes, &*self.state)
1063    }
1064
1065    /// Query a uniquely borrowed world
1066    ///
1067    /// Avoids the cost of the dynamic borrow checking performed by [`query`][Self::query].
1068    pub fn query_mut<'q>(&'q mut self, world: &'q mut World) -> PreparedQueryIter<'q, Q> {
1069        assert_borrow::<Q>();
1070
1071        if self.memo != world.memo() {
1072            *self = Self::prepare(world);
1073        }
1074
1075        let meta = world.entities_meta();
1076        let archetypes = world.archetypes_inner();
1077
1078        let state: &'q [(usize, <Q::Fetch as Fetch<'q>>::State)] =
1079            unsafe { mem::transmute(&*self.state) };
1080
1081        unsafe { PreparedQueryIter::new(meta, archetypes, state.iter()) }
1082    }
1083}
1084
1085/// Combined borrow of a [`PreparedQuery`] and a [`World`]
1086pub struct PreparedQueryBorrow<'q, Q: Query> {
1087    meta: &'q [EntityMeta],
1088    archetypes: &'q [Archetype],
1089    state: &'q [(usize, <Q::Fetch as Fetch<'static>>::State)],
1090}
1091
1092impl<'q, Q: Query> PreparedQueryBorrow<'q, Q> {
1093    fn new(
1094        meta: &'q [EntityMeta],
1095        archetypes: &'q [Archetype],
1096        state: &'q [(usize, <Q::Fetch as Fetch<'static>>::State)],
1097    ) -> Self {
1098        for (idx, state) in state {
1099            Q::Fetch::borrow(&archetypes[*idx], *state);
1100        }
1101
1102        Self {
1103            meta,
1104            archetypes,
1105            state,
1106        }
1107    }
1108
1109    /// Execute the prepared query
1110    // The lifetime narrowing here is required for soundness.
1111    pub fn iter<'i>(&'i mut self) -> PreparedQueryIter<'i, Q> {
1112        let state: &'i [(usize, <Q::Fetch as Fetch<'i>>::State)] =
1113            unsafe { mem::transmute(self.state) };
1114
1115        unsafe { PreparedQueryIter::new(self.meta, self.archetypes, state.iter()) }
1116    }
1117}
1118
1119impl<Q: Query> Drop for PreparedQueryBorrow<'_, Q> {
1120    fn drop(&mut self) {
1121        for (idx, state) in self.state {
1122            Q::Fetch::release(&self.archetypes[*idx], *state);
1123        }
1124    }
1125}
1126
1127/// Iterates over all entities matching a [`PreparedQuery`]
1128pub struct PreparedQueryIter<'q, Q: Query> {
1129    meta: &'q [EntityMeta],
1130    archetypes: &'q [Archetype],
1131    state: SliceIter<'q, (usize, <Q::Fetch as Fetch<'q>>::State)>,
1132    iter: ChunkIter<Q>,
1133}
1134
1135impl<'q, Q: Query> PreparedQueryIter<'q, Q> {
1136    /// # Safety
1137    ///
1138    /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
1139    /// dynamic borrow checks or by representing exclusive access to the `World`.
1140    unsafe fn new(
1141        meta: &'q [EntityMeta],
1142        archetypes: &'q [Archetype],
1143        state: SliceIter<'q, (usize, <Q::Fetch as Fetch<'q>>::State)>,
1144    ) -> Self {
1145        Self {
1146            meta,
1147            archetypes,
1148            state,
1149            iter: ChunkIter::empty(),
1150        }
1151    }
1152}
1153
1154unsafe impl<Q: Query> Send for PreparedQueryIter<'_, Q> {}
1155unsafe impl<Q: Query> Sync for PreparedQueryIter<'_, Q> {}
1156
1157impl<'q, Q: Query> Iterator for PreparedQueryIter<'q, Q> {
1158    type Item = (Entity, QueryItem<'q, Q>);
1159
1160    #[inline(always)]
1161    fn next(&mut self) -> Option<Self::Item> {
1162        loop {
1163            match unsafe { self.iter.next() } {
1164                None => {
1165                    let (idx, state) = self.state.next()?;
1166                    let archetype = &self.archetypes[*idx];
1167                    self.iter = ChunkIter {
1168                        entities: archetype.entities(),
1169                        fetch: Q::Fetch::execute(archetype, *state),
1170                        position: 0,
1171                        len: archetype.len() as usize,
1172                    };
1173                    continue;
1174                }
1175                Some((id, components)) => {
1176                    return Some((
1177                        Entity {
1178                            id,
1179                            generation: unsafe { self.meta.get_unchecked(id as usize).generation },
1180                        },
1181                        components,
1182                    ));
1183                }
1184            }
1185        }
1186    }
1187
1188    fn size_hint(&self) -> (usize, Option<usize>) {
1189        let n = self.len();
1190        (n, Some(n))
1191    }
1192}
1193
1194impl<Q: Query> ExactSizeIterator for PreparedQueryIter<'_, Q> {
1195    fn len(&self) -> usize {
1196        self.state
1197            .clone()
1198            .map(|(idx, _)| self.archetypes[*idx].len() as usize)
1199            .sum::<usize>()
1200            + self.iter.remaining()
1201    }
1202}
1203
1204#[cfg(test)]
1205mod tests {
1206    use super::*;
1207
1208    #[test]
1209    fn access_order() {
1210        assert!(Access::Write > Access::Read);
1211        assert!(Access::Read > Access::Iterate);
1212        assert!(Some(Access::Iterate) > None);
1213    }
1214}
hv_ecs/query.rs

hv_ecs/
query.rs
