update archetypes for run criterias

Author: mockersf

crates/bevy_ecs/src/schedule/executor_parallel.rs

@@ -311,11 +311,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,13 +47,16 @@ pub enum ShouldRun {
 pub(crate) struct BoxedRunCriteria {
     criteria_system: Option<BoxedSystem<(), ShouldRun>>,
     initialized: bool,
+    archetype_generation: ArchetypeGeneration,
 }
 
 impl Default for BoxedRunCriteria {
     fn default() -> Self {
         Self {
             criteria_system: None,
             initialized: false,
+            // MAX ensures access information will be initialized on first run.
+            archetype_generation: ArchetypeGeneration::new(usize::MAX),
         }
     }
 }
@@ -70,6 +73,21 @@ impl BoxedRunCriteria {
                 run_criteria.initialize(world);
                 self.initialized = true;
             }
+            let archetypes = world.archetypes();
+            let old_generation = self.archetype_generation;
+            let new_generation = archetypes.generation();
+            if old_generation != new_generation {
+                let archetype_index_range = if old_generation.value() == usize::MAX {
+                    0..archetypes.len()
+                } else {
+                    old_generation.value()..archetypes.len()
+                };
+                for archetype in archetypes.archetypes[archetype_index_range].iter() {
+                    run_criteria.new_archetype(archetype);
+                }
+
+                self.archetype_generation = new_generation;
+            }
             let should_run = run_criteria.run((), world);
             run_criteria.apply_buffers(world);
             should_run
@@ -93,6 +111,7 @@ pub(crate) struct RunCriteriaContainer {
     pub label: Option<BoxedRunCriteriaLabel>,
     pub before: Vec<BoxedRunCriteriaLabel>,
     pub after: Vec<BoxedRunCriteriaLabel>,
+    archetype_generation: ArchetypeGeneration,
 }
 
 impl RunCriteriaContainer {
@@ -106,6 +125,8 @@ impl RunCriteriaContainer {
             label: descriptor.label,
             before: descriptor.before,
             after: descriptor.after,
+            // MAX ensures access information will be initialized on first run.
+            archetype_generation: ArchetypeGeneration::new(usize::MAX),
         }
     }
 
@@ -122,6 +143,32 @@ impl RunCriteriaContainer {
             RunCriteriaInner::Piped { system, .. } => system.initialize(world),
         }
     }
+
+    pub fn update_archetypes(&mut self, world: &World) {
+        let archetypes = world.archetypes();
+        let old_generation = self.archetype_generation;
+        let new_generation = archetypes.generation();
+        if old_generation == new_generation {
+            return;
+        }
+        let archetype_index_range = if old_generation.value() == usize::MAX {
+            0..archetypes.len()
+        } else {
+            old_generation.value()..archetypes.len()
+        };
+        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<BoxedRunCriteriaLabel> for RunCriteriaContainer {

crates/bevy_ecs/src/schedule/stage.rs

@@ -777,6 +777,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 {
@@ -850,6 +851,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,
@@ -2096,4 +2098,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]);
+    }
 }