Performance tweaks

Author: james-j-obrien

crates/bevy_ecs/Cargo.toml

@@ -29,6 +29,7 @@ rustc-hash = "1.1"
 downcast-rs = "1.2"
 serde = "1"
 thiserror = "1.0"
+smallvec = "1.6"
 
 [dev-dependencies]
 rand = "0.8"

crates/bevy_ecs/src/lib.rs

@@ -53,6 +53,7 @@ pub mod prelude {
             Commands, Deferred, In, IntoSystem, Local, NonSend, NonSendMut, ParallelCommands,
             ParamSet, Query, ReadOnlySystem, Res, ResMut, Resource, System, SystemParamFunction,
         },
+        term_query::{TermQuery, TermQueryState},
         world::{EntityMut, EntityRef, EntityWorldMut, FromWorld, World},
     };
 }

crates/bevy_ecs/src/term_query/builder.rs

@@ -5,18 +5,18 @@ use crate::{
     prelude::{Component, With, Without, World},
 };
 
-use super::{ComponentTerm, QueryTermGroup, Term, TermQueryState};
+use super::{ComponentTerm, QueryTermGroup, Term, TermQueryState, TermVec};
 
 pub struct QueryBuilder<'w, Q: QueryTermGroup = ()> {
-    terms: Vec<Term>,
+    terms: TermVec<Term>,
     current_term: usize,
     world: &'w mut World,
     _marker: PhantomData<Q>,
 }
 
 impl<'w, Q: QueryTermGroup> QueryBuilder<'w, Q> {
     pub fn new(world: &'w mut World) -> Self {
-        let mut terms = Vec::new();
+        let mut terms = TermVec::new();
         Q::init_terms(world, &mut terms);
         Self {
             current_term: terms.len(),

crates/bevy_ecs/src/term_query/iter.rs

@@ -3,22 +3,27 @@ use std::{marker::PhantomData, slice};
 use crate::{
     archetype::{ArchetypeEntity, ArchetypeId, Archetypes},
     component::Tick,
+    entity::Entity,
     query::DebugCheckedUnwrap,
-    storage::Tables,
+    storage::{TableId, TableRow, Tables},
     world::unsafe_world_cell::UnsafeWorldCell,
 };
 
-use super::{Fetchable, FetchedTerm, QueryTermGroup, TermQueryState, TermState};
+use super::{Fetchable, FetchedTerm, QueryTermGroup, TermQueryState, TermState, TermVec};
 
 pub struct TermQueryCursor<'w, 's> {
+    table_id_iter: slice::Iter<'s, TableId>,
     archetype_id_iter: slice::Iter<'s, ArchetypeId>,
+    table_entities: &'w [Entity],
     archetype_entities: &'w [ArchetypeEntity],
-    term_state: Vec<TermState<'w>>,
+    term_state: TermVec<TermState<'w>>,
     current_len: usize,
     current_row: usize,
+    dense: bool,
 }
 
 impl<'w, 's> TermQueryCursor<'w, 's> {
+    #[inline]
     unsafe fn new<Q: QueryTermGroup>(
         world: UnsafeWorldCell<'w>,
         query_state: &'s TermQueryState<Q>,
@@ -27,63 +32,116 @@ impl<'w, 's> TermQueryCursor<'w, 's> {
     ) -> Self {
         let term_state = query_state.init_term_state(world, last_run, this_run);
         Self {
+            table_id_iter: query_state.matched_table_ids.iter(),
             archetype_id_iter: query_state.matched_archetype_ids.iter(),
+            table_entities: &[],
             archetype_entities: &[],
+            dense: term_state.iter().all(|t| t.dense()),
             term_state,
             current_len: 0,
             current_row: 0,
         }
     }
 
+    #[inline(always)]
     unsafe fn next<Q: QueryTermGroup>(
         &mut self,
         tables: &'w Tables,
         archetypes: &'w Archetypes,
         query_state: &'s TermQueryState<Q>,
-    ) -> Option<Vec<FetchedTerm<'w>>> {
-        loop {
-            if self.current_row == self.current_len {
-                let archetype_id = self.archetype_id_iter.next()?;
-                let archetype = archetypes.get(*archetype_id).debug_checked_unwrap();
-                // SAFETY: `archetype` and `tables` are from the world that `fetch/filter` were created for,
-                // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
-                let table = tables.get(archetype.table_id()).debug_checked_unwrap();
-                query_state
+    ) -> Option<TermVec<FetchedTerm<'w>>> {
+        if self.dense {
+            loop {
+                // we are on the beginning of the query, or finished processing a table, so skip to the next
+                if self.current_row == self.current_len {
+                    let table_id = self.table_id_iter.next()?;
+                    let table = tables.get(*table_id).debug_checked_unwrap();
+                    // SAFETY: `table` is from the world that `fetch/filter` were created for,
+                    // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
+                    query_state
+                        .terms
+                        .iter()
+                        .zip(self.term_state.iter_mut())
+                        .for_each(|(term, state)| term.set_table(state, table));
+                    self.table_entities = table.entities();
+                    self.current_len = table.entity_count();
+                    self.current_row = 0;
+                    continue;
+                }
+
+                // SAFETY: set_table was called prior.
+                // `current_row` is a table row in range of the current table, because if it was not, then the if above would have been executed.
+                let entity = *self.table_entities.get_unchecked(self.current_row);
+                let row = TableRow::new(self.current_row);
+
+                // SAFETY:
+                // - set_table was called prior.
+                // - `current_row` must be a table row in range of the current table,
+                //   because if it was not, then the if above would have been executed.
+                // - fetch is only called once for each `entity`.
+                self.current_row += 1;
+
+                if query_state
                     .terms
                     .iter()
                     .zip(self.term_state.iter_mut())
-                    .for_each(|(term, state)| term.set_table(state, table));
-                self.archetype_entities = archetype.entities();
-                self.current_len = archetype.len();
-                self.current_row = 0;
-                continue;
+                    .all(|(term, state)| term.filter_fetch(state, entity, row))
+                {
+                    return Some(
+                        query_state
+                            .terms
+                            .iter()
+                            .zip(self.term_state.iter_mut())
+                            .map(|(term, state)| term.fetch(state, entity, row))
+                            .collect(),
+                    );
+                }
             }
-
-            // SAFETY:
-            // - set_archetype was called prior.
-            // - `current_row` must be an archetype index row in range of the current archetype,
-            //   because if it was not, then the if above would have been executed.
-            // - fetch is only called once for each `archetype_entity`.
-            let archetype_entity = self.archetype_entities.get_unchecked(self.current_row);
-            self.current_row += 1;
-
-            let entity = archetype_entity.entity();
-            let row = archetype_entity.table_row();
-            // Apply filters
-            if query_state
-                .terms
-                .iter()
-                .zip(self.term_state.iter_mut())
-                .all(|(term, state)| term.filter_fetch(state, entity, row))
-            {
-                return Some(
+        } else {
+            loop {
+                if self.current_row == self.current_len {
+                    let archetype_id = self.archetype_id_iter.next()?;
+                    let archetype = archetypes.get(*archetype_id).debug_checked_unwrap();
+                    // SAFETY: `archetype` and `tables` are from the world that `fetch/filter` were created for,
+                    // `fetch_state`/`filter_state` are the states that `fetch/filter` were initialized with
+                    let table = tables.get(archetype.table_id()).debug_checked_unwrap();
                     query_state
                         .terms
                         .iter()
                         .zip(self.term_state.iter_mut())
-                        .map(|(term, state)| term.fetch(state, entity, row))
-                        .collect(),
-                );
+                        .for_each(|(term, state)| term.set_table(state, table));
+                    self.archetype_entities = archetype.entities();
+                    self.current_len = archetype.len();
+                    self.current_row = 0;
+                    continue;
+                }
+
+                // SAFETY:
+                // - set_archetype was called prior.
+                // - `current_row` must be an archetype index row in range of the current archetype,
+                //   because if it was not, then the if above would have been executed.
+                // - fetch is only called once for each `archetype_entity`.
+                let archetype_entity = self.archetype_entities.get_unchecked(self.current_row);
+                self.current_row += 1;
+
+                let entity = archetype_entity.entity();
+                let row = archetype_entity.table_row();
+                // Apply filters
+                if query_state
+                    .terms
+                    .iter()
+                    .zip(self.term_state.iter_mut())
+                    .all(|(term, state)| term.filter_fetch(state, entity, row))
+                {
+                    return Some(
+                        query_state
+                            .terms
+                            .iter()
+                            .zip(self.term_state.iter_mut())
+                            .map(|(term, state)| term.fetch(state, entity, row))
+                            .collect(),
+                    );
+                }
             }
         }
     }
@@ -97,6 +155,7 @@ pub struct TermQueryIterUntyped<'w, 's> {
 }
 
 impl<'w, 's> TermQueryIterUntyped<'w, 's> {
+    #[inline]
     pub unsafe fn new<Q: QueryTermGroup>(
         world: UnsafeWorldCell<'w>,
         query_state: &'s TermQueryState<Q>,
@@ -113,8 +172,9 @@ impl<'w, 's> TermQueryIterUntyped<'w, 's> {
 }
 
 impl<'w, 's> Iterator for TermQueryIterUntyped<'w, 's> {
-    type Item = Vec<FetchedTerm<'w>>;
+    type Item = TermVec<FetchedTerm<'w>>;
 
+    #[inline(always)]
     fn next(&mut self) -> Option<Self::Item> {
         unsafe {
             self.cursor
@@ -129,6 +189,7 @@ pub struct TermQueryIter<'w, 's, Q: QueryTermGroup> {
 }
 
 impl<'w, 's, Q: QueryTermGroup> TermQueryIter<'w, 's, Q> {
+    #[inline(always)]
     pub unsafe fn new(
         world: UnsafeWorldCell<'w>,
         query_state: &'s TermQueryState<Q>,
@@ -145,6 +206,7 @@ impl<'w, 's, Q: QueryTermGroup> TermQueryIter<'w, 's, Q> {
 impl<'w, 's, Q: QueryTermGroup> Iterator for TermQueryIter<'w, 's, Q> {
     type Item = Q::Item<'w>;
 
+    #[inline(always)]
     fn next(&mut self) -> Option<Self::Item> {
         unsafe {
             self.inner

crates/bevy_ecs/src/term_query/mod.rs

@@ -349,4 +349,40 @@ mod tests {
         assert_eq!(Some(&B(0)), entity.get::<B>());
         assert_eq!(None, entity.get::<C>());
     }
+
+    #[derive(Component)]
+    #[component(storage = "SparseSet")]
+    struct S(usize);
+
+    #[test]
+    fn term_query_sparse_set() {
+        let mut world = World::new();
+        let entity_a = world.spawn((A(0), S(1))).id();
+
+        let mut query = world.term_query::<(Entity, &A, &S)>();
+
+        let (e, a, s) = query.single(&world);
+        assert_eq!(entity_a, e);
+        assert_eq!(0, a.0);
+        assert_eq!(1, s.0);
+    }
+
+    #[test]
+    fn term_query_iteration() {
+        let mut world = World::new();
+        let entity = world.spawn((A(1), B(0), C(0))).id();
+        world.spawn_batch((1..1000).map(|i| (A(i), B(0))));
+
+        let mut query = world.term_query::<(&A, &mut B)>();
+
+        query
+            .iter_mut(&mut world)
+            .for_each(|(a, mut b)| b.0 = a.0 * 2);
+
+        let mut query = world.term_query_filtered::<(Entity, &B), With<C>>();
+        let (e, b) = query.single(&world);
+
+        assert_eq!(e, entity);
+        assert_eq!(2, b.0);
+    }
 }

crates/bevy_ecs/src/term_query/query.rs

@@ -102,6 +102,7 @@ impl<'w, 's, Q: QueryTermGroup, F: QueryTermGroup> TermQuery<'w, 's, Q, F> {
         self.get_single_mut().unwrap()
     }
 
+    #[inline]
     pub fn get_single_mut(&mut self) -> Result<Q::Item<'_>, QuerySingleError> {
         // SAFETY:
         // the query ensures mutable access to the components it accesses, and the query