ecs: move copy of QueryBorrow into bevy_ecs and fix perf regressions

This commit is contained in:
Carter Anderson 2020-07-28 16:17:21 -07:00
parent 77f4e60c8c
commit 64cc382477
4 changed files with 322 additions and 96 deletions

View file

@ -40,7 +40,8 @@ pub struct World {
entities: Entities,
index: HashMap<Vec<TypeId>, u32>,
removed_components: HashMap<TypeId, Vec<Entity>>,
archetypes: Vec<Archetype>,
#[allow(missing_docs)]
pub archetypes: Vec<Archetype>,
archetype_generation: u64,
}

View file

@ -1,12 +1,10 @@
pub use super::Query;
use super::TypeAccess;
use crate::{
resource::{FetchResource, ResourceQuery, Resources, UnsafeClone},
system::{ArchetypeAccess, Commands, System, SystemId, ThreadLocalExecution},
};
use core::marker::PhantomData;
use hecs::{
Component, ComponentError, Entity, Fetch, Query as HecsQuery, QueryBorrow, Ref, RefMut, World,
};
use hecs::{Fetch, Query as HecsQuery, World};
use std::borrow::Cow;
pub struct SystemFn<State, F, ThreadLocalF, Init, SetArchetypeAccess>
@ -133,92 +131,6 @@ macro_rules! impl_into_foreach_system {
};
}
pub struct Query<'a, Q: HecsQuery> {
world: &'a World,
archetype_access: &'a ArchetypeAccess,
_marker: PhantomData<Q>,
}
#[derive(Debug)]
pub enum QueryComponentError {
CannotReadArchetype,
CannotWriteArchetype,
ComponentError(ComponentError),
}
impl<'a, Q: HecsQuery> Query<'a, Q> {
pub fn iter(&mut self) -> QueryBorrow<'_, Q> {
self.world.query::<Q>()
}
/// Gets a reference to the entity's component of the given type. This will fail if the entity does not have
/// the given component type or if the given component type does not match this query.
pub fn get<T: Component>(&self, entity: Entity) -> Result<Ref<'_, T>, QueryComponentError> {
if let Some(location) = self.world.get_entity_location(entity) {
if self
.archetype_access
.immutable
.contains(location.archetype as usize)
|| self
.archetype_access
.mutable
.contains(location.archetype as usize)
{
self.world
.get(entity)
.map_err(|err| QueryComponentError::ComponentError(err))
} else {
Err(QueryComponentError::CannotReadArchetype)
}
} else {
Err(QueryComponentError::ComponentError(
ComponentError::NoSuchEntity,
))
}
}
/// Gets a mutable reference to the entity's component of the given type. This will fail if the entity does not have
/// the given component type or if the given component type does not match this query.
pub fn get_mut<T: Component>(
&self,
entity: Entity,
) -> Result<RefMut<'_, T>, QueryComponentError> {
if let Some(location) = self.world.get_entity_location(entity) {
if self
.archetype_access
.mutable
.contains(location.archetype as usize)
{
self.world
.get_mut(entity)
.map_err(|err| QueryComponentError::ComponentError(err))
} else {
Err(QueryComponentError::CannotWriteArchetype)
}
} else {
Err(QueryComponentError::ComponentError(
ComponentError::NoSuchEntity,
))
}
}
pub fn removed<C: Component>(&self) -> &[Entity] {
self.world.removed::<C>()
}
/// Sets the entity's component to the given value. This will fail if the entity does not already have
/// the given component type or if the given component type does not match this query.
pub fn set<T: Component>(
&self,
entity: Entity,
component: T,
) -> Result<(), QueryComponentError> {
let mut current = self.get_mut::<T>(entity)?;
*current = component;
Ok(())
}
}
struct QuerySystemState {
archetype_accesses: Vec<ArchetypeAccess>,
commands: Commands,
@ -265,11 +177,7 @@ macro_rules! impl_into_query_system {
let ($($resource,)*) = resources.query_system::<($($resource,)*)>(id);
let mut i = 0;
$(
let $query = Query::<$query> {
world,
archetype_access: &state.archetype_accesses[i],
_marker: PhantomData::default(),
};
let $query = Query::<$query>::new(world, &state.archetype_accesses[i]);
i += 1;
)*

View file

@ -3,9 +3,11 @@ mod into_system;
#[cfg(feature = "profiler")]
mod profiler;
mod system;
mod query;
pub use commands::*;
pub use into_system::*;
#[cfg(feature = "profiler")]
pub use profiler::*;
pub use system::*;
pub use query::*;

View file

@ -0,0 +1,315 @@
use hecs::{
Component, ComponentError, Entity, Query as HecsQuery, Ref, RefMut, World, Archetype, Access, Without, With, Fetch,
};
use std::marker::PhantomData;
use crate::ArchetypeAccess;
pub struct Query<'a, Q: HecsQuery> {
pub(crate) world: &'a World,
pub(crate) archetype_access: &'a ArchetypeAccess,
_marker: PhantomData<Q>,
}
#[derive(Debug)]
pub enum QueryComponentError {
CannotReadArchetype,
CannotWriteArchetype,
ComponentError(ComponentError),
}
impl<'a, Q: HecsQuery> Query<'a, Q> {
#[inline]
pub fn new(world: &'a World, archetype_access: &'a ArchetypeAccess) -> Self {
Self {
world,
archetype_access,
_marker: PhantomData::default()
}
}
#[inline]
pub fn iter(&mut self) -> QueryBorrow<'_, Q> {
QueryBorrow::new(&self.world.archetypes)
}
/// Gets a reference to the entity's component of the given type. This will fail if the entity does not have
/// the given component type or if the given component type does not match this query.
pub fn get<T: Component>(&self, entity: Entity) -> Result<Ref<'_, T>, QueryComponentError> {
if let Some(location) = self.world.get_entity_location(entity) {
if self
.archetype_access
.immutable
.contains(location.archetype as usize)
|| self
.archetype_access
.mutable
.contains(location.archetype as usize)
{
self.world
.get(entity)
.map_err(|err| QueryComponentError::ComponentError(err))
} else {
Err(QueryComponentError::CannotReadArchetype)
}
} else {
Err(QueryComponentError::ComponentError(
ComponentError::NoSuchEntity,
))
}
}
/// Gets a mutable reference to the entity's component of the given type. This will fail if the entity does not have
/// the given component type or if the given component type does not match this query.
pub fn get_mut<T: Component>(
&self,
entity: Entity,
) -> Result<RefMut<'_, T>, QueryComponentError> {
if let Some(location) = self.world.get_entity_location(entity) {
if self
.archetype_access
.mutable
.contains(location.archetype as usize)
{
self.world
.get_mut(entity)
.map_err(|err| QueryComponentError::ComponentError(err))
} else {
Err(QueryComponentError::CannotWriteArchetype)
}
} else {
Err(QueryComponentError::ComponentError(
ComponentError::NoSuchEntity,
))
}
}
pub fn removed<C: Component>(&self) -> &[Entity] {
self.world.removed::<C>()
}
/// Sets the entity's component to the given value. This will fail if the entity does not already have
/// the given component type or if the given component type does not match this query.
pub fn set<T: Component>(
&self,
entity: Entity,
component: T,
) -> Result<(), QueryComponentError> {
let mut current = self.get_mut::<T>(entity)?;
*current = component;
Ok(())
}
}
/// A borrow of a `World` sufficient to execute the query `Q`
///
/// Note that borrows are not released until this object is dropped.
pub struct QueryBorrow<'w, Q: HecsQuery> {
archetypes: &'w [Archetype],
borrowed: bool,
_marker: PhantomData<Q>,
}
impl<'w, Q: HecsQuery> QueryBorrow<'w, Q> {
pub(crate) fn new(archetypes: &'w [Archetype]) -> Self {
Self {
archetypes,
borrowed: false,
_marker: PhantomData,
}
}
/// Execute the query
///
/// Must be called only once per query.
#[inline]
pub fn iter<'q>(&'q mut self) -> QueryIter<'q, 'w, Q> {
self.borrow();
QueryIter {
borrow: self,
archetype_index: 0,
iter: None,
}
}
fn borrow(&mut self) {
if self.borrowed {
panic!(
"called QueryBorrow::iter twice on the same borrow; construct a new query instead"
);
}
for x in self.archetypes {
// TODO: Release prior borrows on failure?
if Q::Fetch::access(x) >= Some(Access::Read) {
Q::Fetch::borrow(x);
}
}
self.borrowed = true;
}
/// Transform the query into one that requires a certain component without borrowing it
///
/// This can be useful when the component needs to be borrowed elsewhere and it isn't necessary
/// for the iterator to expose its data directly.
///
/// Equivalent to using a query type wrapped in `With`.
///
/// # Example
/// ```
/// # use hecs::*;
/// let mut world = World::new();
/// let a = world.spawn((123, true, "abc"));
/// let b = world.spawn((456, false));
/// let c = world.spawn((42, "def"));
/// let entities = world.query::<(Entity, &i32)>()
/// .with::<bool>()
/// .iter()
/// .map(|(e, &i)| (e, i)) // Copy out of the world
/// .collect::<Vec<_>>();
/// assert!(entities.contains(&(a, 123)));
/// assert!(entities.contains(&(b, 456)));
/// ```
pub fn with<T: Component>(self) -> QueryBorrow<'w, With<T, Q>> {
self.transform()
}
/// Transform the query into one that skips entities having a certain component
///
/// Equivalent to using a query type wrapped in `Without`.
///
/// # Example
/// ```
/// # use hecs::*;
/// let mut world = World::new();
/// let a = world.spawn((123, true, "abc"));
/// let b = world.spawn((456, false));
/// let c = world.spawn((42, "def"));
/// let entities = world.query::<(Entity, &i32)>()
/// .without::<bool>()
/// .iter()
/// .map(|(e, &i)| (e, i)) // Copy out of the world
/// .collect::<Vec<_>>();
/// assert_eq!(entities, &[(c, 42)]);
/// ```
pub fn without<T: Component>(self) -> QueryBorrow<'w, Without<T, Q>> {
self.transform()
}
/// Helper to change the type of the query
fn transform<R: HecsQuery>(mut self) -> QueryBorrow<'w, R> {
let x = QueryBorrow {
archetypes: self.archetypes,
borrowed: self.borrowed,
_marker: PhantomData,
};
// Ensure `Drop` won't fire redundantly
self.borrowed = false;
x
}
}
unsafe impl<'w, Q: HecsQuery> Send for QueryBorrow<'w, Q> {}
unsafe impl<'w, Q: HecsQuery> Sync for QueryBorrow<'w, Q> {}
impl<'w, Q: HecsQuery> Drop for QueryBorrow<'w, Q> {
fn drop(&mut self) {
if self.borrowed {
for x in self.archetypes {
if Q::Fetch::access(x) >= Some(Access::Read) {
Q::Fetch::release(x);
}
}
}
}
}
impl<'q, 'w, Q: HecsQuery> IntoIterator for &'q mut QueryBorrow<'w, Q> {
type IntoIter = QueryIter<'q, 'w, Q>;
type Item = <Q::Fetch as Fetch<'q>>::Item;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
/// Iterator over the set of entities with the components in `Q`
pub struct QueryIter<'q, 'w, Q: HecsQuery> {
borrow: &'q mut QueryBorrow<'w, Q>,
archetype_index: u32,
iter: Option<ChunkIter<Q>>,
}
unsafe impl<'q, 'w, Q: HecsQuery> Send for QueryIter<'q, 'w, Q> {}
unsafe impl<'q, 'w, Q: HecsQuery> Sync for QueryIter<'q, 'w, Q> {}
impl<'q, 'w, Q: HecsQuery> Iterator for QueryIter<'q, 'w, Q> {
type Item = <Q::Fetch as Fetch<'q>>::Item;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
loop {
match self.iter {
None => {
let archetype = self.borrow.archetypes.get(self.archetype_index as usize)?;
self.archetype_index += 1;
unsafe {
self.iter = Q::Fetch::get(archetype, 0).map(|fetch| ChunkIter {
fetch,
len: archetype.len(),
});
}
}
Some(ref mut iter) => match unsafe { iter.next() } {
None => {
self.iter = None;
continue;
}
Some(components) => {
return Some(components);
}
},
}
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let n = self.len();
(n, Some(n))
}
}
impl<'q, 'w, Q: HecsQuery> ExactSizeIterator for QueryIter<'q, 'w, Q> {
fn len(&self) -> usize {
self.borrow
.archetypes
.iter()
.filter(|&x| Q::Fetch::access(x).is_some())
.map(|x| x.len() as usize)
.sum()
}
}
struct ChunkIter<Q: HecsQuery> {
fetch: Q::Fetch,
len: u32,
}
impl<Q: HecsQuery> ChunkIter<Q> {
#[inline]
unsafe fn next<'a, 'w>(&mut self) -> Option<<Q::Fetch as Fetch<'a>>::Item> {
loop {
if self.len == 0 {
return None;
}
self.len -= 1;
if self.fetch.should_skip() {
// we still need to progress the iterator
let _ = self.fetch.next();
continue;
}
break Some(self.fetch.next());
}
}
}