diff --git a/crates/bevy_ecs/hecs/src/world.rs b/crates/bevy_ecs/hecs/src/world.rs index eb5ed6e01f..2f69f56877 100644 --- a/crates/bevy_ecs/hecs/src/world.rs +++ b/crates/bevy_ecs/hecs/src/world.rs @@ -40,7 +40,8 @@ pub struct World { entities: Entities, index: HashMap, u32>, removed_components: HashMap>, - archetypes: Vec, + #[allow(missing_docs)] + pub archetypes: Vec, archetype_generation: u64, } diff --git a/crates/bevy_ecs/src/system/into_system.rs b/crates/bevy_ecs/src/system/into_system.rs index 167b8f58f3..2fec7ae9e7 100644 --- a/crates/bevy_ecs/src/system/into_system.rs +++ b/crates/bevy_ecs/src/system/into_system.rs @@ -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 @@ -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, -} - -#[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::() - } - - /// 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(&self, entity: Entity) -> Result, 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( - &self, - entity: Entity, - ) -> Result, 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(&self) -> &[Entity] { - self.world.removed::() - } - - /// 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( - &self, - entity: Entity, - component: T, - ) -> Result<(), QueryComponentError> { - let mut current = self.get_mut::(entity)?; - *current = component; - Ok(()) - } -} - struct QuerySystemState { archetype_accesses: Vec, 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; )* diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index 555b48bc02..ea15404f59 100644 --- a/crates/bevy_ecs/src/system/mod.rs +++ b/crates/bevy_ecs/src/system/mod.rs @@ -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::*; diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs new file mode 100644 index 0000000000..73b18ccb88 --- /dev/null +++ b/crates/bevy_ecs/src/system/query.rs @@ -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, +} + +#[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(&self, entity: Entity) -> Result, 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( + &self, + entity: Entity, + ) -> Result, 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(&self) -> &[Entity] { + self.world.removed::() + } + + /// 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( + &self, + entity: Entity, + component: T, + ) -> Result<(), QueryComponentError> { + let mut current = self.get_mut::(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, +} + +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::() + /// .iter() + /// .map(|(e, &i)| (e, i)) // Copy out of the world + /// .collect::>(); + /// assert!(entities.contains(&(a, 123))); + /// assert!(entities.contains(&(b, 456))); + /// ``` + pub fn with(self) -> QueryBorrow<'w, With> { + 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::() + /// .iter() + /// .map(|(e, &i)| (e, i)) // Copy out of the world + /// .collect::>(); + /// assert_eq!(entities, &[(c, 42)]); + /// ``` + pub fn without(self) -> QueryBorrow<'w, Without> { + self.transform() + } + + /// Helper to change the type of the query + fn transform(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 = >::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>, +} + +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 = >::Item; + + #[inline] + fn next(&mut self) -> Option { + 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) { + 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 { + fetch: Q::Fetch, + len: u32, +} + +impl ChunkIter { + #[inline] + unsafe fn next<'a, 'w>(&mut self) -> Option<>::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()); + } + } +} \ No newline at end of file