From a7bab755eeb70eb9ced00647aa944f8240f9ac19 Mon Sep 17 00:00:00 2001 From: Carter Anderson Date: Tue, 14 Jul 2020 19:05:39 -0700 Subject: [PATCH] ecs: add query get safety checks --- crates/bevy_ecs/hecs/src/world.rs | 5 + crates/bevy_ecs/src/into_system.rs | 229 ++++++++++++++---- crates/bevy_ecs/src/parallel_executor.rs | 3 + crates/bevy_ecs/src/schedule.rs | 1 + crates/bevy_ecs/src/system.rs | 49 +++- .../src/camera/visible_entities.rs | 14 +- crates/bevy_ui/src/ui_update_system.rs | 8 +- 7 files changed, 248 insertions(+), 61 deletions(-) diff --git a/crates/bevy_ecs/hecs/src/world.rs b/crates/bevy_ecs/hecs/src/world.rs index d4b6bfb3b2..f31f9ac77d 100644 --- a/crates/bevy_ecs/hecs/src/world.rs +++ b/crates/bevy_ecs/hecs/src/world.rs @@ -554,6 +554,11 @@ impl World { pub fn archetypes_generation(&self) -> ArchetypesGeneration { ArchetypesGeneration(self.archetype_generation) } + + /// Retrieves the entity's current location, if it exists + pub fn get_entity_location(&self, entity: Entity) -> Option { + self.entities.get(entity).ok() + } } unsafe impl Send for World {} diff --git a/crates/bevy_ecs/src/into_system.rs b/crates/bevy_ecs/src/into_system.rs index e09c8c11b5..dea494bf10 100644 --- a/crates/bevy_ecs/src/into_system.rs +++ b/crates/bevy_ecs/src/into_system.rs @@ -9,13 +9,15 @@ use hecs::{ }; use std::borrow::Cow; -pub struct SystemFn +pub struct SystemFn where - F: FnMut(&World, &Resources) + Send + Sync, - ThreadLocalF: FnMut(&mut World, &mut Resources) + Send + Sync, + F: FnMut(&World, &Resources, &ArchetypeAccess, &mut State) + Send + Sync, + ThreadLocalF: FnMut(&mut World, &mut Resources, &mut State) + Send + Sync, Init: FnMut(&mut Resources) + Send + Sync, - SetArchetypeAccess: FnMut(&World, &mut ArchetypeAccess) + Send + Sync, + SetArchetypeAccess: FnMut(&World, &mut ArchetypeAccess, &mut State) + Send + Sync, + State: Send + Sync, { + pub state: State, pub func: F, pub thread_local_func: ThreadLocalF, pub init_func: Init, @@ -26,20 +28,21 @@ where pub set_archetype_access: SetArchetypeAccess, } -impl System - for SystemFn +impl System + for SystemFn where - F: FnMut(&World, &Resources) + Send + Sync, - ThreadLocalF: FnMut(&mut World, &mut Resources) + Send + Sync, + F: FnMut(&World, &Resources, &ArchetypeAccess, &mut State) + Send + Sync, + ThreadLocalF: FnMut(&mut World, &mut Resources, &mut State) + Send + Sync, Init: FnMut(&mut Resources) + Send + Sync, - SetArchetypeAccess: FnMut(&World, &mut ArchetypeAccess) + Send + Sync, + SetArchetypeAccess: FnMut(&World, &mut ArchetypeAccess, &mut State) + Send + Sync, + State: Send + Sync, { fn name(&self) -> Cow<'static, str> { self.name.clone() } fn update_archetype_access(&mut self, world: &World) { - (self.set_archetype_access)(world, &mut self.archetype_access); + (self.set_archetype_access)(world, &mut self.archetype_access, &mut self.state); } fn get_archetype_access(&self) -> &ArchetypeAccess { @@ -51,11 +54,11 @@ where } fn run(&mut self, world: &World, resources: &Resources) { - (self.func)(world, resources); + (self.func)(world, resources, &self.archetype_access, &mut self.state); } fn run_thread_local(&mut self, world: &mut World, resources: &mut Resources) { - (self.thread_local_func)(world, resources); + (self.thread_local_func)(world, resources, &mut self.state); } fn initialize(&mut self, resources: &mut Resources) { @@ -90,33 +93,32 @@ macro_rules! impl_into_foreach_system { #[allow(unused_unsafe)] fn system(mut self) -> Box { let id = SystemId::new(); - let commands = Commands::default(); - let thread_local_commands = commands.clone(); Box::new(SystemFn { + state: Commands::default(), thread_local_execution: ThreadLocalExecution::NextFlush, name: core::any::type_name::().into(), id, - func: move |world, resources| { + func: move |world, resources, _archetype_access, state| { <<($($resource,)*) as ResourceQuery>::Fetch as FetchResource>::borrow(&resources.resource_archetypes); { let ($($resource,)*) = resources.query_system::<($($resource,)*)>(id); + let commands = state.clone(); for ($($component,)*) in world.query::<($($component,)*)>().iter() { fn_call!(self, ($($commands, commands)*), ($($resource),*), ($($component),*)) } } <<($($resource,)*) as ResourceQuery>::Fetch as FetchResource>::release(&resources.resource_archetypes); }, - thread_local_func: move |world, resources| { - thread_local_commands.apply(world, resources); + thread_local_func: move |world, resources, state| { + state.apply(world, resources); }, init_func: move |resources| { <($($resource,)*)>::initialize(resources, Some(id)); }, archetype_access: ArchetypeAccess::default(), - set_archetype_access: |world, archetype_access| { - for archetype in world.archetypes() { - archetype_access.set_access_for_query::<($($component,)*)>(world); - } + set_archetype_access: |world, archetype_access, _state| { + archetype_access.clear(); + archetype_access.set_access_for_query::<($($component,)*)>(world); }, }) } @@ -126,25 +128,71 @@ 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::() } - pub fn get(&self, entity: Entity) -> Result, ComponentError> { - // TODO: Check if request matches query - self.world.get(entity) + 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.world + .get(entity) + .map_err(|err| QueryComponentError::ComponentError(err)) + } else { + Err(QueryComponentError::CannotReadArchetype) + } + } else { + Err(QueryComponentError::ComponentError( + ComponentError::NoSuchEntity, + )) + } } - pub fn get_mut(&self, entity: Entity) -> Result, ComponentError> { - // TODO: Check if request matches query - self.world.get_mut(entity) + 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, + )) + } } } +struct QuerySystemState { + archetype_accesses: Vec, + commands: Commands, +} + pub trait IntoQuerySystem { fn system(self) -> Box; } @@ -165,38 +213,58 @@ macro_rules! impl_into_query_system { #[allow(non_snake_case)] #[allow(unused_variables)] #[allow(unused_unsafe)] + #[allow(unused_assignments)] + #[allow(unused_mut)] fn system(mut self) -> Box { let id = SystemId::new(); - let commands = Commands::default(); - let thread_local_commands = commands.clone(); + $(let $query = ArchetypeAccess::default();)* Box::new(SystemFn { + state: QuerySystemState { + archetype_accesses: vec![ + $($query,)* + ], + commands: Commands::default(), + }, thread_local_execution: ThreadLocalExecution::NextFlush, id, name: core::any::type_name::().into(), - func: move |world, resources| { + func: move |world, resources, archetype_access, state| { <<($($resource,)*) as ResourceQuery>::Fetch as FetchResource>::borrow(&resources.resource_archetypes); { let ($($resource,)*) = resources.query_system::<($($resource,)*)>(id); - $(let $query = Query::<$query> { - world, - _marker: PhantomData::default(), - };)* + let mut i = 0; + $( + let $query = Query::<$query> { + world, + archetype_access: &state.archetype_accesses[i], + _marker: PhantomData::default(), + }; + i += 1; + )* + let commands = state.commands.clone(); fn_call!(self, ($($commands, commands)*), ($($resource),*), ($($query),*)) } <<($($resource,)*) as ResourceQuery>::Fetch as FetchResource>::release(&resources.resource_archetypes); }, - thread_local_func: move |world, resources| { - thread_local_commands.apply(world, resources); + thread_local_func: move |world, resources, state| { + state.commands.apply(world, resources); }, init_func: move |resources| { <($($resource,)*)>::initialize(resources, Some(id)); }, archetype_access: ArchetypeAccess::default(), - set_archetype_access: |world, archetype_access| { - for archetype in world.archetypes() { - $(archetype_access.set_access_for_query::<$query>(world);)* - } + set_archetype_access: |world, archetype_access, state| { + archetype_access.clear(); + let mut i = 0; + let mut access: &mut ArchetypeAccess; + $( + access = &mut state.archetype_accesses[i]; + access.clear(); + access.set_access_for_query::<$query>(world); + archetype_access.union(access); + i += 1; + )* }, }) } @@ -300,12 +368,13 @@ where { fn thread_local_system(mut self) -> Box { Box::new(SystemFn { - thread_local_func: move |world, resources| { + state: (), + thread_local_func: move |world, resources, _| { self.run(world, resources); }, - func: |_, _| {}, + func: |_, _, _, _| {}, init_func: |_| {}, - set_archetype_access: |_, _| {}, + set_archetype_access: |_, _, _| {}, thread_local_execution: ThreadLocalExecution::Immediate, name: core::any::type_name::().into(), id: SystemId::new(), @@ -326,3 +395,77 @@ where self(world, resources); } } + +#[cfg(test)] +mod tests { + use super::IntoQuerySystem; + use crate::{Query, ResMut, Resources, Schedule}; + use hecs::{Entity, With, World}; + + struct A; + struct B; + struct C; + struct D; + + #[test] + fn query_system_gets() { + fn query_system( + mut ran: ResMut, + mut entity_query: Query>, + b_query: Query<&B>, + a_c_query: Query<(&A, &C)>, + d_query: Query<&D>, + ) { + let entities = entity_query.iter().iter().collect::>(); + assert!( + b_query.get::(entities[0]).is_err(), + "entity 0 should not have B" + ); + assert!( + b_query.get::(entities[1]).is_ok(), + "entity 1 should have B" + ); + assert!( + b_query.get::(entities[1]).is_ok(), + "entity 1 should have A, and it should (unintuitively) be accessible from b_query because b_query grabs read access to the (A,B) archetype"); + assert!( + b_query.get::(entities[3]).is_err(), + "entity 3 should have D, but it shouldn't be accessible from b_query" + ); + assert!( + b_query.get::(entities[2]).is_err(), + "entity 2 has C, but it shouldn't be accessible from b_query" + ); + assert!( + a_c_query.get::(entities[2]).is_ok(), + "entity 2 has C, and it should be accessible from a_c_query" + ); + assert!( + a_c_query.get::(entities[3]).is_err(), + "entity 3 should have D, but it shouldn't be accessible from b_query" + ); + assert!( + d_query.get::(entities[3]).is_ok(), + "entity 3 should have D" + ); + + *ran = true; + } + + let mut world = World::default(); + let mut resources = Resources::default(); + resources.insert(false); + world.spawn((A,)); + world.spawn((A, B)); + world.spawn((A, C)); + world.spawn((A, D)); + + let mut schedule = Schedule::default(); + schedule.add_stage("update"); + schedule.add_system_to_stage("update", query_system.system()); + + schedule.run(&mut world, &mut resources); + + assert!(*resources.get::().unwrap(), "system ran"); + } +} diff --git a/crates/bevy_ecs/src/parallel_executor.rs b/crates/bevy_ecs/src/parallel_executor.rs index 8488f3bdbc..d5771d03fd 100644 --- a/crates/bevy_ecs/src/parallel_executor.rs +++ b/crates/bevy_ecs/src/parallel_executor.rs @@ -226,6 +226,9 @@ impl ExecutorStage { system.run_thread_local(world, resources); self.finished_systems.insert(thread_local_index); self.sender.send(thread_local_index).unwrap(); + + // TODO: if archetype generation has changed, call "prepare" on all systems after this one + run_ready_result = RunReadyResult::Ok; } else { // wait for a system to finish, then run its dependents diff --git a/crates/bevy_ecs/src/schedule.rs b/crates/bevy_ecs/src/schedule.rs index 73842f08de..24c173d2af 100644 --- a/crates/bevy_ecs/src/schedule.rs +++ b/crates/bevy_ecs/src/schedule.rs @@ -103,6 +103,7 @@ impl Schedule { #[cfg(feature = "profiler")] crate::profiler::profiler_start(resources, system.name().clone()); let mut system = system.lock().unwrap(); + system.update_archetype_access(world); match system.thread_local_execution() { ThreadLocalExecution::NextFlush => system.run(world, resources), ThreadLocalExecution::Immediate => { diff --git a/crates/bevy_ecs/src/system.rs b/crates/bevy_ecs/src/system.rs index abeac2ff37..6c3659eb6f 100644 --- a/crates/bevy_ecs/src/system.rs +++ b/crates/bevy_ecs/src/system.rs @@ -1,7 +1,7 @@ use crate::{Resources, World}; -use std::borrow::Cow; use fixedbitset::FixedBitSet; -use hecs::{Query, Access}; +use hecs::{Access, Query}; +use std::borrow::Cow; #[derive(Copy, Clone)] pub enum ThreadLocalExecution { @@ -52,8 +52,6 @@ impl ArchetypeAccess { where Q: Query, { - self.immutable.clear(); - self.mutable.clear(); let iterator = world.archetypes(); let bits = iterator.len(); self.immutable.grow(bits); @@ -67,4 +65,45 @@ impl ArchetypeAccess { Access::Iterate => (), }); } -} \ No newline at end of file + + pub fn clear(&mut self) { + self.immutable.clear(); + self.mutable.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::ArchetypeAccess; + use hecs::World; + + struct A; + struct B; + struct C; + + #[test] + fn query_archetype_access() { + let mut world = World::default(); + let e1 = world.spawn((A,)); + let e2 = world.spawn((A, B)); + let e3 = world.spawn((A, B, C)); + + let mut access = ArchetypeAccess::default(); + access.set_access_for_query::<(&A,)>(&world); + + let e1_archetype = world.get_entity_location(e1).unwrap().archetype as usize; + let e2_archetype = world.get_entity_location(e2).unwrap().archetype as usize; + let e3_archetype = world.get_entity_location(e3).unwrap().archetype as usize; + + assert!(access.immutable.contains(e1_archetype)); + assert!(access.immutable.contains(e2_archetype)); + assert!(access.immutable.contains(e3_archetype)); + + let mut access = ArchetypeAccess::default(); + access.set_access_for_query::<(&A, &B)>(&world); + + assert!(access.immutable.contains(e1_archetype) == false); + assert!(access.immutable.contains(e2_archetype)); + assert!(access.immutable.contains(e3_archetype)); + } +} diff --git a/crates/bevy_render/src/camera/visible_entities.rs b/crates/bevy_render/src/camera/visible_entities.rs index 59a8c12ff2..01c392ec91 100644 --- a/crates/bevy_render/src/camera/visible_entities.rs +++ b/crates/bevy_render/src/camera/visible_entities.rs @@ -21,24 +21,22 @@ impl VisibleEntities { } pub fn visible_entities_system( - mut camera_query: Query<(Entity, &Camera, &mut VisibleEntities)>, - mut entities_query: Query<(Entity, &Draw)>, - transform_query: Query<&Transform>, - _transform_entities_query: Query<(&Draw, &Transform)>, // ensures we can optionally access Transforms + mut camera_query: Query<(&Camera, &Transform, &mut VisibleEntities)>, + mut draw_query: Query<(Entity, &Draw)>, + draw_transform_query: Query<(&Draw, &Transform)>, ) { - for (camera_entity, _camera, visible_entities) in &mut camera_query.iter() { + for (_camera, camera_transform, visible_entities) in &mut camera_query.iter() { visible_entities.value.clear(); - let camera_transform = transform_query.get::(camera_entity).unwrap(); let camera_position = camera_transform.value.w_axis().truncate(); let mut no_transform_order = 0.0; let mut transparent_entities = Vec::new(); - for (entity, draw) in &mut entities_query.iter() { + for (entity, draw) in &mut draw_query.iter() { if !draw.is_visible { continue; } - let order = if let Ok(transform) = transform_query.get::(entity) { + let order = if let Ok(transform) = draw_transform_query.get::(entity) { let position = transform.value.w_axis().truncate(); // smaller distances are sorted to lower indices by using the distance from the camera FloatOrd((camera_position - position).length()) diff --git a/crates/bevy_ui/src/ui_update_system.rs b/crates/bevy_ui/src/ui_update_system.rs index 6de8eb664b..3ec87c7cc4 100644 --- a/crates/bevy_ui/src/ui_update_system.rs +++ b/crates/bevy_ui/src/ui_update_system.rs @@ -1,6 +1,6 @@ use super::Node; use bevy_core::transform::run_on_hierarchy; -use bevy_ecs::{Entity, Query, Res}; +use bevy_ecs::{Entity, Query, Res, Without}; use bevy_transform::prelude::{Children, Parent, Translation}; use bevy_window::Windows; use glam::Vec2; @@ -15,8 +15,8 @@ pub struct Rect { pub fn ui_update_system( windows: Res, + mut orphan_node_query: Query>, mut node_query: Query<(Entity, &mut Node, &mut Translation)>, - parent_query: Query<&Parent>, children_query: Query<&Children>, ) { let window_size = if let Some(window) = windows.get_primary() { @@ -24,11 +24,9 @@ pub fn ui_update_system( } else { return; }; - let orphan_nodes = node_query + let orphan_nodes = orphan_node_query .iter() .iter() - // TODO: replace this filter with a legion query filter (when SimpleQuery gets support for filters) - .filter(|(entity, _, _)| parent_query.get::(*entity).is_err()) .map(|(e, _, _)| e) .collect::>(); let mut window_rect = Rect {