Optimize iteration

src/archetype.rs

@@ -104,9 +104,7 @@ impl Archetype {
         assert_eq!(self.types[state].id, TypeId::of::<T>());
 
         unsafe {
-            NonNull::new_unchecked(
-                self.data.get_unchecked(state).storage.as_ptr().cast::<T>() as *mut T
-            )
+            NonNull::new_unchecked(self.data.get_unchecked(state).storage.as_ptr().cast::<T>())
         }
     }
 

src/lib.rs

@@ -41,18 +41,27 @@ extern crate std;
 
 extern crate alloc;
 
+macro_rules! reverse_apply {
+    ($m: ident [] $($reversed:tt)*) => {
+        $m!{$($reversed),*}  // base case
+    };
+    ($m: ident [$first:tt $($rest:tt)*] $($reversed:tt)*) => {
+        reverse_apply!{$m [$($rest)*] $first $($reversed)*}
+    };
+}
+
 /// Imagine macro parameters, but more like those Russian dolls.
 ///
-/// Calls m!(A, B, C), m!(A, B), m!(B), and m!() for i.e. (m, A, B, C)
+/// Calls m!(), m!(A), m!(A, B), and m!(A, B, C) for i.e. (m, A, B, C)
 /// where m is any macro, for any number of parameters.
 macro_rules! smaller_tuples_too {
-    ($m: ident, $ty: ident) => {
+    ($m: ident, $next: tt) => {
         $m!{}
-        $m!{$ty}
+        $m!{$next}
     };
-    ($m: ident, $ty: ident, $($tt: ident),*) => {
-        smaller_tuples_too!{$m, $($tt),*}
-        $m!{$ty, $($tt),*}
+    ($m: ident, $next: tt, $($rest: tt),*) => {
+        smaller_tuples_too!{$m, $($rest),*}
+        reverse_apply!{$m [$next $($rest)*]}
     };
 }
 

src/query.rs

@@ -604,17 +604,15 @@ impl<T> Clone for FetchSatisfies<T> {
 ///
 /// Note that borrows are not released until this object is dropped.
 pub struct QueryBorrow<'w, Q: Query> {
-    meta: &'w [EntityMeta],
-    archetypes: &'w [Archetype],
+    world: &'w World,
     borrowed: bool,
     _marker: PhantomData<Q>,
 }
 
 impl<'w, Q: Query> QueryBorrow<'w, Q> {
-    pub(crate) fn new(meta: &'w [EntityMeta], archetypes: &'w [Archetype]) -> Self {
+    pub(crate) fn new(world: &'w World) -> Self {
         Self {
-            meta,
-            archetypes,
+            world,
             borrowed: false,
             _marker: PhantomData,
         }
@@ -624,13 +622,13 @@ impl<'w, Q: Query> QueryBorrow<'w, Q> {
     // The lifetime narrowing here is required for soundness.
     pub fn iter(&mut self) -> QueryIter<'_, Q> {
         self.borrow();
-        unsafe { QueryIter::new(self.meta, self.archetypes.iter()) }
+        unsafe { QueryIter::new(self.world) }
     }
 
     /// Provide random access to the query results
     pub fn view(&mut self) -> View<'_, Q> {
         self.borrow();
-        unsafe { View::new(self.meta, self.archetypes) }
+        unsafe { View::new(self.world.entities_meta(), self.world.archetypes_inner()) }
     }
 
     /// Like `iter`, but returns child iterators of at most `batch_size` elements
@@ -639,14 +637,20 @@ impl<'w, Q: Query> QueryBorrow<'w, Q> {
     // The lifetime narrowing here is required for soundness.
     pub fn iter_batched(&mut self, batch_size: u32) -> BatchedIter<'_, Q> {
         self.borrow();
-        unsafe { BatchedIter::new(self.meta, self.archetypes.iter(), batch_size) }
+        unsafe {
+            BatchedIter::new(
+                self.world.entities_meta(),
+                self.world.archetypes_inner().iter(),
+                batch_size,
+            )
+        }
     }
 
     fn borrow(&mut self) {
         if self.borrowed {
             return;
         }
-        for x in self.archetypes {
+        for x in self.world.archetypes() {
             if x.is_empty() {
                 continue;
             }
@@ -709,8 +713,7 @@ impl<'w, Q: Query> QueryBorrow<'w, Q> {
     /// Helper to change the type of the query
     fn transform<R: Query>(mut self) -> QueryBorrow<'w, R> {
         let x = QueryBorrow {
-            meta: self.meta,
-            archetypes: self.archetypes,
+            world: self.world,
             borrowed: self.borrowed,
             _marker: PhantomData,
         };
@@ -726,7 +729,7 @@ unsafe impl<'w, Q: Query> Sync for QueryBorrow<'w, Q> where for<'a> Q::Item<'a>:
 impl<'w, Q: Query> Drop for QueryBorrow<'w, Q> {
     fn drop(&mut self) {
         if self.borrowed {
-            for x in self.archetypes {
+            for x in self.world.archetypes() {
                 if x.is_empty() {
                     continue;
                 }
@@ -749,8 +752,8 @@ impl<'q, 'w, Q: Query> IntoIterator for &'q mut QueryBorrow<'w, Q> {
 
 /// Iterator over the set of entities with the components in `Q`
 pub struct QueryIter<'q, Q: Query> {
-    meta: &'q [EntityMeta],
-    archetypes: SliceIter<'q, Archetype>,
+    world: &'q World,
+    archetypes: core::ops::Range<usize>,
     iter: ChunkIter<Q>,
 }
 
@@ -759,13 +762,26 @@ impl<'q, Q: Query> QueryIter<'q, Q> {
     ///
     /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
     /// dynamic borrow checks or by representing exclusive access to the `World`.
-    unsafe fn new(meta: &'q [EntityMeta], archetypes: SliceIter<'q, Archetype>) -> Self {
+    unsafe fn new(world: &'q World) -> Self {
+        let n = world.archetypes().len();
         Self {
-            meta,
-            archetypes,
+            world,
+            archetypes: 0..n,
             iter: ChunkIter::empty(),
         }
     }
+
+    /// Advance query to the next archetype
+    ///
+    /// Outlined from `Iterator::next` for improved iteration performance.
+    fn next_archetype(&mut self) -> Option<()> {
+        let archetype = self.archetypes.next()?;
+        let archetype = unsafe { self.world.archetypes_inner().get_unchecked(archetype) };
+        let state = Q::Fetch::prepare(archetype);
+        let fetch = state.map(|state| Q::Fetch::execute(archetype, state));
+        self.iter = fetch.map_or(ChunkIter::empty(), |fetch| ChunkIter::new(archetype, fetch));
+        Some(())
+    }
 }
 
 unsafe impl<'q, Q: Query> Send for QueryIter<'q, Q> where for<'a> Q::Item<'a>: Send {}
@@ -779,22 +795,19 @@ impl<'q, Q: Query> Iterator for QueryIter<'q, Q> {
         loop {
             match unsafe { self.iter.next() } {
                 None => {
-                    let archetype = self.archetypes.next()?;
-                    let state = Q::Fetch::prepare(archetype);
-                    let fetch = state.map(|state| Q::Fetch::execute(archetype, state));
-                    self.iter = fetch.map_or(ChunkIter::empty(), |fetch| ChunkIter {
-                        entities: archetype.entities(),
-                        fetch,
-                        position: 0,
-                        len: archetype.len() as usize,
-                    });
+                    self.next_archetype()?;
                     continue;
                 }
                 Some((id, components)) => {
                     return Some((
                         Entity {
                             id,
-                            generation: unsafe { self.meta.get_unchecked(id as usize).generation },
+                            generation: unsafe {
+                                self.world
+                                    .entities_meta()
+                                    .get_unchecked(id as usize)
+                                    .generation
+                            },
                         },
                         components,
                     ));
@@ -813,6 +826,7 @@ impl<'q, Q: Query> ExactSizeIterator for QueryIter<'q, Q> {
     fn len(&self) -> usize {
         self.archetypes
             .clone()
+            .map(|x| unsafe { self.world.archetypes_inner().get_unchecked(x) })
             .filter(|&x| Q::Fetch::access(x).is_some())
             .map(|x| x.len() as usize)
             .sum::<usize>()
@@ -826,17 +840,22 @@ pub struct QueryMut<'q, Q: Query> {
 }
 
 impl<'q, Q: Query> QueryMut<'q, Q> {
-    pub(crate) fn new(meta: &'q [EntityMeta], archetypes: &'q mut [Archetype]) -> Self {
+    pub(crate) fn new(world: &'q mut World) -> Self {
         assert_borrow::<Q>();
 
         Self {
-            iter: unsafe { QueryIter::new(meta, archetypes.iter()) },
+            iter: unsafe { QueryIter::new(world) },
         }
     }
 
     /// Provide random access to the query results
     pub fn view(&mut self) -> View<'_, Q> {
-        unsafe { View::new(self.iter.meta, self.iter.archetypes.as_slice()) }
+        unsafe {
+            View::new(
+                self.iter.world.entities_meta(),
+                self.iter.world.archetypes_inner(),
+            )
+        }
     }
 
     /// Transform the query into one that requires another query be satisfied
@@ -856,15 +875,21 @@ impl<'q, Q: Query> QueryMut<'q, Q> {
     /// Helper to change the type of the query
     fn transform<R: Query>(self) -> QueryMut<'q, R> {
         QueryMut {
-            iter: unsafe { QueryIter::new(self.iter.meta, self.iter.archetypes) },
+            iter: unsafe { QueryIter::new(self.iter.world) },
         }
     }
 
     /// Like `into_iter`, but returns child iterators of at most `batch_size` elements
     ///
     /// Useful for distributing work over a threadpool.
     pub fn into_iter_batched(self, batch_size: u32) -> BatchedIter<'q, Q> {
-        unsafe { BatchedIter::new(self.iter.meta, self.iter.archetypes, batch_size) }
+        unsafe {
+            BatchedIter::new(
+                self.iter.world.entities_meta(),
+                self.iter.world.archetypes_inner().iter(),
+                batch_size,
+            )
+        }
     }
 }
 
@@ -904,6 +929,15 @@ struct ChunkIter<Q: Query> {
 }
 
 impl<Q: Query> ChunkIter<Q> {
+    fn new(archetype: &Archetype, fetch: Q::Fetch) -> Self {
+        Self {
+            entities: archetype.entities(),
+            fetch,
+            position: 0,
+            len: archetype.len() as usize,
+        }
+    }
+
     fn empty() -> Self {
         Self {
             entities: NonNull::dangling(),
@@ -978,14 +1012,12 @@ impl<'q, Q: Query> Iterator for BatchedIter<'q, Q> {
             let fetch = state.map(|state| Q::Fetch::execute(archetype, state));
             if let Some(fetch) = fetch {
                 self.batch += 1;
+                let mut state = ChunkIter::new(archetype, fetch);
+                state.position = offset as usize;
+                state.len = (offset + self.batch_size.min(archetype.len() - offset)) as usize;
                 return Some(Batch {
                     meta: self.meta,
-                    state: ChunkIter {
-                        entities: archetype.entities(),
-                        fetch,
-                        len: (offset + self.batch_size.min(archetype.len() - offset)) as usize,
-                        position: offset as usize,
-                    },
+                    state,
                 });
             } else {
                 self.archetypes = archetypes;
@@ -1048,10 +1080,12 @@ macro_rules! tuple_impl {
                 $($name::borrow(archetype, $name);)*
             }
             #[allow(unused_variables)]
+            #[cold]
             fn prepare(archetype: &Archetype) -> Option<Self::State> {
                 Some(($($name::prepare(archetype)?,)*))
             }
             #[allow(unused_variables, non_snake_case, clippy::unused_unit)]
+            #[inline(always)]
             fn execute(archetype: &Archetype, state: Self::State) -> Self {
                 let ($($name,)*) = state;
                 ($($name::execute(archetype, $name),)*)
@@ -1278,12 +1312,7 @@ impl<'q, Q: Query> Iterator for PreparedQueryIter<'q, Q> {
                 None => {
                     let (idx, state) = self.state.next()?;
                     let archetype = &self.archetypes[*idx];
-                    self.iter = ChunkIter {
-                        entities: archetype.entities(),
-                        fetch: Q::Fetch::execute(archetype, *state),
-                        position: 0,
-                        len: archetype.len() as usize,
-                    };
+                    self.iter = ChunkIter::new(archetype, Q::Fetch::execute(archetype, *state));
                     continue;
                 }
                 Some((id, components)) => {
@@ -1463,12 +1492,9 @@ impl<'a, Q: Query> Iterator for ViewIter<'a, Q> {
                 None => {
                     let archetype = self.archetypes.next()?;
                     let fetch = self.fetches.next()?;
-                    self.iter = fetch.clone().map_or(ChunkIter::empty(), |fetch| ChunkIter {
-                        entities: archetype.entities(),
-                        fetch,
-                        position: 0,
-                        len: archetype.len() as usize,
-                    });
+                    self.iter = fetch
+                        .clone()
+                        .map_or(ChunkIter::empty(), |fetch| ChunkIter::new(archetype, fetch));
                     continue;
                 }
                 Some((id, components)) => {

src/world.rs

@@ -396,7 +396,7 @@ impl World {
     /// assert!(entities.contains(&(b, 456, false)));
     /// ```
     pub fn query<Q: Query>(&self) -> QueryBorrow<'_, Q> {
-        QueryBorrow::new(&self.entities.meta, &self.archetypes.archetypes)
+        QueryBorrow::new(self)
     }
 
     /// Query a uniquely borrowed world
@@ -405,17 +405,19 @@ impl World {
     /// that, unlike [`query`](Self::query), this returns an `IntoIterator` which can be passed
     /// directly to a `for` loop.
     pub fn query_mut<Q: Query>(&mut self) -> QueryMut<'_, Q> {
-        QueryMut::new(&self.entities.meta, &mut self.archetypes.archetypes)
+        QueryMut::new(self)
     }
 
     pub(crate) fn memo(&self) -> (u64, u32) {
         (self.id, self.archetypes.generation())
     }
 
+    #[inline(always)]
     pub(crate) fn entities_meta(&self) -> &[EntityMeta] {
         &self.entities.meta
     }
 
+    #[inline(always)]
     pub(crate) fn archetypes_inner(&self) -> &[Archetype] {
         &self.archetypes.archetypes
     }
@@ -822,6 +824,7 @@ impl World {
     ///
     /// Useful for dynamically scheduling concurrent queries by checking borrows in advance, and for
     /// efficient serialization.
+    #[inline(always)]
     pub fn archetypes(&self) -> impl ExactSizeIterator<Item = &'_ Archetype> + '_ {
         self.archetypes_inner().iter()
     }