update archetypes for run criterias (#2177)

fixes #2000 

archetypes were not updated for run criteria on a stage or on a system set

crates/bevy_ecs/src/schedule/executor_parallel.rs

@@ -301,11 +301,7 @@ impl ParallelExecutor {
 
     #[cfg(test)]
     fn emit_event(&self, event: SchedulingEvent) {
-        self.events_sender
-            .as_ref()
-            .unwrap()
-            .try_send(event)
-            .unwrap();
+        let _ = self.events_sender.as_ref().unwrap().try_send(event);
     }
 }
 

crates/bevy_ecs/src/schedule/run_criteria.rs

@@ -1,5 +1,5 @@
 use crate::{
-    archetype::{Archetype, ArchetypeComponentId},
+    archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration},
     component::ComponentId,
     query::Access,
     schedule::{BoxedRunCriteriaLabel, GraphNode, RunCriteriaLabel},
@@ -47,6 +47,7 @@ pub enum ShouldRun {
 pub(crate) struct BoxedRunCriteria {
     criteria_system: Option<BoxedSystem<(), ShouldRun>>,
     initialized: bool,
+    archetype_generation: ArchetypeGeneration,
 }
 
 impl Default for BoxedRunCriteria {
@@ -54,6 +55,7 @@ impl Default for BoxedRunCriteria {
         Self {
             criteria_system: None,
             initialized: false,
+            archetype_generation: ArchetypeGeneration::initial(),
         }
     }
 }
@@ -70,6 +72,15 @@ impl BoxedRunCriteria {
                 run_criteria.initialize(world);
                 self.initialized = true;
             }
+            let archetypes = world.archetypes();
+            let new_generation = archetypes.generation();
+            let old_generation = std::mem::replace(&mut self.archetype_generation, new_generation);
+            let archetype_index_range = old_generation.value()..new_generation.value();
+
+            for archetype in archetypes.archetypes[archetype_index_range].iter() {
+                run_criteria.new_archetype(archetype);
+            }
+
             let should_run = run_criteria.run((), world);
             run_criteria.apply_buffers(world);
             should_run
@@ -93,6 +104,7 @@ pub(crate) struct RunCriteriaContainer {
     pub label: Option<BoxedRunCriteriaLabel>,
     pub before: Vec<BoxedRunCriteriaLabel>,
     pub after: Vec<BoxedRunCriteriaLabel>,
+    archetype_generation: ArchetypeGeneration,
 }
 
 impl RunCriteriaContainer {
@@ -106,6 +118,7 @@ impl RunCriteriaContainer {
             label: descriptor.label,
             before: descriptor.before,
             after: descriptor.after,
+            archetype_generation: ArchetypeGeneration::initial(),
         }
     }
 
@@ -122,6 +135,25 @@ impl RunCriteriaContainer {
             RunCriteriaInner::Piped { system, .. } => system.initialize(world),
         }
     }
+
+    pub fn update_archetypes(&mut self, world: &World) {
+        let archetypes = world.archetypes();
+        let new_generation = archetypes.generation();
+        let old_generation = std::mem::replace(&mut self.archetype_generation, new_generation);
+        let archetype_index_range = old_generation.value()..new_generation.value();
+        for archetype in archetypes.archetypes[archetype_index_range].iter() {
+            match &mut self.inner {
+                RunCriteriaInner::Single(system) => {
+                    system.new_archetype(archetype);
+                }
+
+                RunCriteriaInner::Piped { system, .. } => {
+                    system.new_archetype(archetype);
+                }
+            }
+        }
+        self.archetype_generation = new_generation;
+    }
 }
 
 impl GraphNode for RunCriteriaContainer {

crates/bevy_ecs/src/schedule/stage.rs

@@ -775,6 +775,7 @@ impl Stage for SystemStage {
             for index in 0..self.run_criteria.len() {
                 let (run_criteria, tail) = self.run_criteria.split_at_mut(index);
                 let mut criteria = &mut tail[0];
+                criteria.update_archetypes(world);
                 match &mut criteria.inner {
                     RunCriteriaInner::Single(system) => criteria.should_run = system.run((), world),
                     RunCriteriaInner::Piped {
@@ -848,6 +849,7 @@ impl Stage for SystemStage {
                 for index in 0..run_criteria.len() {
                     let (run_criteria, tail) = run_criteria.split_at_mut(index);
                     let criteria = &mut tail[0];
+                    criteria.update_archetypes(world);
                     match criteria.should_run {
                         ShouldRun::No => (),
                         ShouldRun::Yes => criteria.should_run = ShouldRun::No,
@@ -2094,4 +2096,77 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn run_criteria_with_query() {
+        struct Foo;
+
+        fn even_number_of_entities_critiera(query: Query<&Foo>) -> ShouldRun {
+            if query.iter().len() % 2 == 0 {
+                ShouldRun::Yes
+            } else {
+                ShouldRun::No
+            }
+        }
+
+        fn spawn_entity(mut commands: crate::prelude::Commands) {
+            commands.spawn().insert(Foo);
+        }
+
+        fn count_entities(query: Query<&Foo>, mut res: ResMut<Vec<usize>>) {
+            res.push(query.iter().len());
+        }
+
+        let mut world = World::new();
+        world.insert_resource(Vec::<usize>::new());
+        let mut stage = SystemStage::parallel()
+            .with_system(spawn_entity.system().label("spawn"))
+            .with_system_set(
+                SystemSet::new()
+                    .with_run_criteria(even_number_of_entities_critiera.system())
+                    .with_system(count_entities.system().before("spawn")),
+            );
+        stage.run(&mut world);
+        stage.run(&mut world);
+        stage.run(&mut world);
+        stage.run(&mut world);
+        assert_eq!(*world.get_resource::<Vec<usize>>().unwrap(), vec![0, 2]);
+    }
+
+    #[test]
+    fn stage_run_criteria_with_query() {
+        struct Foo;
+
+        fn even_number_of_entities_critiera(query: Query<&Foo>) -> ShouldRun {
+            if query.iter().len() % 2 == 0 {
+                ShouldRun::Yes
+            } else {
+                ShouldRun::No
+            }
+        }
+
+        fn spawn_entity(mut commands: crate::prelude::Commands) {
+            commands.spawn().insert(Foo);
+        }
+
+        fn count_entities(query: Query<&Foo>, mut res: ResMut<Vec<usize>>) {
+            res.push(query.iter().len());
+        }
+
+        let mut world = World::new();
+        world.insert_resource(Vec::<usize>::new());
+        let mut stage_spawn = SystemStage::parallel().with_system(spawn_entity.system());
+        let mut stage_count = SystemStage::parallel()
+            .with_run_criteria(even_number_of_entities_critiera.system())
+            .with_system(count_entities.system());
+        stage_count.run(&mut world);
+        stage_spawn.run(&mut world);
+        stage_count.run(&mut world);
+        stage_spawn.run(&mut world);
+        stage_count.run(&mut world);
+        stage_spawn.run(&mut world);
+        stage_count.run(&mut world);
+        stage_spawn.run(&mut world);
+        assert_eq!(*world.get_resource::<Vec<usize>>().unwrap(), vec![0, 2]);
+    }
 }

crates/bevy_ecs/src/system/exclusive_system.rs

@@ -1,4 +1,5 @@
 use crate::{
+    archetype::ArchetypeGeneration,
     system::{check_system_change_tick, BoxedSystem, IntoSystem, SystemId},
     world::World,
 };
@@ -74,6 +75,7 @@ where
 
 pub struct ExclusiveSystemCoerced {
     system: BoxedSystem<(), ()>,
+    archetype_generation: ArchetypeGeneration,
 }
 
 impl ExclusiveSystem for ExclusiveSystemCoerced {
@@ -86,6 +88,15 @@ impl ExclusiveSystem for ExclusiveSystemCoerced {
     }
 
     fn run(&mut self, world: &mut World) {
+        let archetypes = world.archetypes();
+        let new_generation = archetypes.generation();
+        let old_generation = std::mem::replace(&mut self.archetype_generation, new_generation);
+        let archetype_index_range = old_generation.value()..new_generation.value();
+
+        for archetype in archetypes.archetypes[archetype_index_range].iter() {
+            self.system.new_archetype(archetype);
+        }
+
         self.system.run((), world);
         self.system.apply_buffers(world);
     }
@@ -106,6 +117,7 @@ where
     fn exclusive_system(self) -> ExclusiveSystemCoerced {
         ExclusiveSystemCoerced {
             system: Box::new(self.system()),
+            archetype_generation: ArchetypeGeneration::initial(),
         }
     }
 }
@@ -148,4 +160,26 @@ mod tests {
         stage.run(&mut world);
         assert_eq!(*world.get_resource::<usize>().unwrap(), 1);
     }
+
+    #[test]
+    fn update_archetype_for_exclusive_system_coerced() {
+        struct Foo;
+
+        fn spawn_entity(mut commands: crate::prelude::Commands) {
+            commands.spawn().insert(Foo);
+        }
+
+        fn count_entities(query: Query<&Foo>, mut res: ResMut<Vec<usize>>) {
+            res.push(query.iter().len());
+        }
+
+        let mut world = World::new();
+        world.insert_resource(Vec::<usize>::new());
+        let mut stage = SystemStage::parallel()
+            .with_system(spawn_entity.system())
+            .with_system(count_entities.exclusive_system());
+        stage.run(&mut world);
+        stage.run(&mut world);
+        assert_eq!(*world.get_resource::<Vec<usize>>().unwrap(), vec![0, 1]);
+    }
 }