ecs: add query get safety checks

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<Location> {
+        self.entities.get(entity).ok()
+    }
 }
 
 unsafe impl Send for World {}

crates/bevy_ecs/src/into_system.rs

@@ -9,13 +9,15 @@ use hecs::{
 };
 use std::borrow::Cow;
 
-pub struct SystemFn<F, ThreadLocalF, Init, SetArchetypeAccess>
+pub struct SystemFn<State, F, ThreadLocalF, Init, SetArchetypeAccess>
 where
-    F: FnMut(&World, &Resources) + Send + Sync,
+    F: FnMut(&World, &Resources, &ArchetypeAccess, &mut State) + Send + Sync,
-    ThreadLocalF: FnMut(&mut World, &mut Resources) + 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<F, ThreadLocalF, Init, SetArchetypeAccess> System
+impl<State, F, ThreadLocalF, Init, SetArchetypeAccess> System
-    for SystemFn<F, ThreadLocalF, Init, SetArchetypeAccess>
+    for SystemFn<State, F, ThreadLocalF, Init, SetArchetypeAccess>
 where
-    F: FnMut(&World, &Resources) + Send + Sync,
+    F: FnMut(&World, &Resources, &ArchetypeAccess, &mut State) + Send + Sync,
-    ThreadLocalF: FnMut(&mut World, &mut Resources) + 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<dyn System> {
                 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::<Self>().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_func: move |world, resources, state| {
-                        thread_local_commands.apply(world, resources);
+                        state.apply(world, resources);
                     },
                     init_func: move |resources| {
                         <($($resource,)*)>::initialize(resources, Some(id));
                     },
                     archetype_access: ArchetypeAccess::default(),
-                    set_archetype_access: |world, archetype_access| {
+                    set_archetype_access: |world, archetype_access, _state| {
-                        for archetype in world.archetypes() {
+                        archetype_access.clear();
                         archetype_access.set_access_for_query::<($($component,)*)>(world);
-                        }
                     },
                 })
             }
@@ -126,23 +128,69 @@ 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>()
     }
 
-    pub fn get<T: Component>(&self, entity: Entity) -> Result<Ref<'_, T>, ComponentError> {
+    pub fn get<T: Component>(&self, entity: Entity) -> Result<Ref<'_, T>, QueryComponentError> {
-        // TODO: Check if request matches query
+        if let Some(location) = self.world.get_entity_location(entity) {
-        self.world.get(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<T: Component>(&self, entity: Entity) -> Result<RefMut<'_, T>, ComponentError> {
+    pub fn get_mut<T: Component>(
-        // TODO: Check if request matches query
+        &self,
-        self.world.get_mut(entity)
+        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,
+            ))
+        }
+    }
+}
+
+struct QuerySystemState {
+    archetype_accesses: Vec<ArchetypeAccess>,
+    commands: Commands,
 }
 
 pub trait IntoQuerySystem<Commands, R, Q> {
@@ -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<dyn System> {
                 let id = SystemId::new();
-                let commands = Commands::default();
+                $(let $query = ArchetypeAccess::default();)*
-                let thread_local_commands = commands.clone();
                 Box::new(SystemFn {
+                    state: QuerySystemState {
+                        archetype_accesses: vec![
+                            $($query,)*
+                        ],
+                        commands: Commands::default(),
+                    },
                     thread_local_execution: ThreadLocalExecution::NextFlush,
                     id,
                     name: core::any::type_name::<Self>().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> {
+                            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_func: move |world, resources, state| {
-                        thread_local_commands.apply(world, resources);
+                        state.commands.apply(world, resources);
                     },
                     init_func: move |resources| {
                         <($($resource,)*)>::initialize(resources, Some(id));
                     },
                     archetype_access: ArchetypeAccess::default(),
-                    set_archetype_access: |world, archetype_access| {
+                    set_archetype_access: |world, archetype_access, state| {
-                        for archetype in world.archetypes() {
+                        archetype_access.clear();
-                           $(archetype_access.set_access_for_query::<$query>(world);)*
+                        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<dyn System> {
         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::<F>().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<bool>,
+            mut entity_query: Query<With<A, Entity>>,
+            b_query: Query<&B>,
+            a_c_query: Query<(&A, &C)>,
+            d_query: Query<&D>,
+        ) {
+            let entities = entity_query.iter().iter().collect::<Vec<Entity>>();
+            assert!(
+                b_query.get::<B>(entities[0]).is_err(),
+                "entity 0 should not have B"
+            );
+            assert!(
+                b_query.get::<B>(entities[1]).is_ok(),
+                "entity 1 should have B"
+            );
+            assert!(
+                b_query.get::<A>(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::<D>(entities[3]).is_err(),
+                "entity 3 should have D, but it shouldn't be accessible from b_query"
+            );
+            assert!(
+                b_query.get::<C>(entities[2]).is_err(),
+                "entity 2 has C, but it shouldn't be accessible from b_query"
+            );
+            assert!(
+                a_c_query.get::<C>(entities[2]).is_ok(),
+                "entity 2 has C, and it should be accessible from a_c_query"
+            );
+            assert!(
+                a_c_query.get::<D>(entities[3]).is_err(),
+                "entity 3 should have D, but it shouldn't be accessible from b_query"
+            );
+            assert!(
+                d_query.get::<D>(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::<bool>().unwrap(), "system ran");
+    }
+}

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

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 => {

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 => (),
             });
     }
+
+    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));
+    }
 }

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 camera_query: Query<(&Camera, &Transform, &mut VisibleEntities)>,
-    mut entities_query: Query<(Entity, &Draw)>,
+    mut draw_query: Query<(Entity, &Draw)>,
-    transform_query: Query<&Transform>,
+    draw_transform_query: Query<(&Draw, &Transform)>,
-    _transform_entities_query: Query<(&Draw, &Transform)>, // ensures we can optionally access Transforms
 ) {
-    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::<Transform>(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::<Transform>(entity) {
+            let order = if let Ok(transform) = draw_transform_query.get::<Transform>(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())

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<Windows>,
+    mut orphan_node_query: Query<Without<Parent, (Entity, &mut Node, &mut Translation)>>,
     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::<Parent>(*entity).is_err())
         .map(|(e, _, _)| e)
         .collect::<Vec<Entity>>();
     let mut window_rect = Rect {