Remove systems from schedule

crates/bevy_ecs/src/schedule/error.rs

@@ -259,3 +259,17 @@ impl ScheduleBuildWarning {
         }
     }
 }
+
+/// Error returned from some `Schedule` methods
+#[derive(Error, Debug)]
+pub enum ScheduleError {
+    /// Operation cannot be completed because the schedule has changed and `Schedule::initialize` needs to be called
+    #[error("Operation cannot be completed because the schedule has changed and `Schedule::initialize` needs to be called")]
+    Uninitialized,
+    /// Method could not find key
+    #[error("Not Found")]
+    NotFound,
+    /// Error initializing schedule
+    #[error("{0}")]
+    ScheduleBuildError(ScheduleBuildError),
+}

crates/bevy_ecs/src/schedule/node.rs

@@ -493,8 +493,8 @@ impl Systems {
     /// # Panics
     ///
     /// If the system with the given key does not exist in this container.
-    pub(crate) fn node_mut(&mut self, key: SystemKey) -> &mut SystemNode {
-        &mut self.nodes[key]
+    pub(crate) fn node_mut(&mut self, key: SystemKey) -> Option<&mut SystemNode> {
+        self.nodes.get_mut(key)
     }
 
     /// Returns `true` if the system with the given key has conditions.
@@ -554,6 +554,25 @@ impl Systems {
         key
     }
 
+    /// Remove a system with [`SystemKey`]
+    pub(crate) fn remove(&mut self, key: SystemKey) -> bool {
+        let mut found = false;
+        if self.nodes.remove(key).is_some() {
+            found = true;
+        }
+
+        if self.conditions.remove(key).is_some() {
+            found = true;
+        }
+
+        if let Some(index) = self.uninit.iter().position(|value| *value == key) {
+            self.uninit.remove(index);
+            found = true;
+        }
+
+        found
+    }
+
     /// Returns `true` if all systems in this container have been initialized.
     pub fn is_initialized(&self) -> bool {
         self.uninit.is_empty()
@@ -646,6 +665,11 @@ impl SystemSets {
         self.sets.get(key).map(|set| &**set)
     }
 
+    /// Returns the kkey for the given system set, returns None if it does not exist.
+    pub fn get_key(&self, set: InternedSystemSet) -> Option<SystemSetKey> {
+        self.ids.get(&set).copied()
+    }
+
     /// Returns the key for the given system set, inserting it into this
     /// container if it does not already exist.
     pub fn get_key_or_insert(&mut self, set: InternedSystemSet) -> SystemSetKey {
@@ -716,6 +740,16 @@ impl SystemSets {
         key
     }
 
+    /// Remove a set with a [`SystemSetKey`]
+    pub(crate) fn remove(&mut self, key: SystemSetKey) -> bool {
+        self.sets.remove(key);
+        self.conditions.remove(key);
+        if let Some(index) = self.uninit.iter().position(|uninit| uninit.key == key) {
+            self.uninit.remove(index);
+        }
+        true
+    }
+
     /// Returns `true` if all system sets' conditions in this container have
     /// been initialized.
     pub fn is_initialized(&self) -> bool {

crates/bevy_ecs/src/schedule/schedule.rs

@@ -379,6 +379,41 @@ impl Schedule {
         self
     }
 
+    /// Removes all systems in a [`SystemSet`]. This will cause the schedule to be rebuilt when
+    /// the schedule is run again. A [`ScheduleError`] is returned if the schedule needs to be
+    /// [`Schedule::initialize`]'d or the `set` is not found.
+    ///
+    /// Note that this can remove all systems of a type if you pass
+    /// the system to this function as systems implicitly create a set based
+    /// on the system type.
+    ///
+    /// ## Example
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # fn my_system() {}
+    /// #
+    /// let mut schedule = Schedule::default();
+    /// // add the system to the schedule
+    /// schedule.add_systems(my_system);
+    /// let mut world = World::default();
+    ///
+    /// // remove the system
+    /// schedule.remove_systems_in_set(my_system, &mut world);
+    /// ```
+    pub fn remove_systems_in_set<M>(
+        &mut self,
+        set: impl IntoSystemSet<M>,
+        world: &mut World,
+    ) -> Result<usize, ScheduleError> {
+        if self.graph.changed {
+            if let Err(e) = self.initialize(world) {
+                return Err(ScheduleError::ScheduleBuildError(e));
+            }
+        }
+        self.graph.remove_systems_in_set(set)
+    }
+
     /// Suppress warnings and errors that would result from systems in these sets having ambiguities
     /// (conflicting access but indeterminate order) with systems in `set`.
     #[track_caller]
@@ -683,6 +718,8 @@ pub struct ScheduleGraph {
     hierarchy: Dag<NodeId>,
     /// Directed acyclic graph of the dependency (which systems/sets have to run before which other systems/sets)
     dependency: Dag<NodeId>,
+    /// Map of systems in each set
+    set_systems: HashMap<SystemSetKey, Vec<SystemKey>>,
     ambiguous_with: UnGraph<NodeId>,
     /// Nodes that are allowed to have ambiguous ordering relationship with any other systems.
     pub ambiguous_with_all: HashSet<NodeId>,
@@ -701,6 +738,7 @@ impl ScheduleGraph {
             system_sets: SystemSets::default(),
             hierarchy: Dag::new(),
             dependency: Dag::new(),
+            set_systems: HashMap::new(),
             ambiguous_with: UnGraph::default(),
             ambiguous_with_all: HashSet::default(),
             conflicting_systems: Vec::new(),
@@ -899,6 +937,58 @@ impl ScheduleGraph {
         AnonymousSet::new(id)
     }
 
+    /// Returns iterator over all [`SystemKey`]'s in a [`SystemSet`]
+    /// Returns `ScheduleBuildError::Uninitialized` if schedule has been changed and `Self::initialize`
+    /// has not been called.
+    pub fn systems_in_set(
+        &self,
+        system_set: InternedSystemSet,
+    ) -> Result<Vec<SystemKey>, ScheduleError> {
+        if self.changed {
+            return Err(ScheduleError::Uninitialized);
+        }
+        let system_set_id = self
+            .system_sets
+            .get_key(system_set)
+            .ok_or(ScheduleError::NotFound)?;
+        self.set_systems
+            .get(&system_set_id)
+            .cloned()
+            .ok_or(ScheduleError::NotFound)
+    }
+
+    /// Remove all systems in a set and any dependencies on those systems and set.
+    pub fn remove_systems_in_set<M>(
+        &mut self,
+        system_set: impl IntoSystemSet<M>,
+    ) -> Result<usize, ScheduleError> {
+        let set = system_set.into_system_set();
+        let interned = set.intern();
+        let keys = self.systems_in_set(interned)?;
+
+        self.changed = true;
+        for &key in &keys {
+            self.systems.remove(key);
+
+            self.hierarchy.graph.remove_node(key.into());
+            self.dependency.graph.remove_node(key.into());
+            self.ambiguous_with.remove_node(key.into());
+            self.ambiguous_with_all.remove(&NodeId::from(key));
+        }
+
+        let Some(set_key) = self.system_sets.get_key(interned) else {
+            return Ok(keys.len());
+        };
+        self.system_sets.remove(set_key);
+        self.set_systems.remove(&set_key);
+        self.hierarchy.graph.remove_node(set_key.into());
+        self.dependency.graph.remove_node(set_key.into());
+        self.ambiguous_with.remove_node(set_key.into());
+        self.ambiguous_with_all.remove(&NodeId::from(set_key));
+
+        Ok(keys.len())
+    }
+
     /// Update the internal graphs (hierarchy, dependency, ambiguity) by adding a single [`GraphInfo`]
     fn update_graphs(&mut self, id: NodeId, graph_info: GraphInfo) {
         self.changed = true;
@@ -1036,6 +1126,7 @@ impl ScheduleGraph {
 
         // flatten: combine `in_set` with `ambiguous_with` information
         let ambiguous_with_flattened = self.get_ambiguous_with_flattened(&set_systems);
+        self.set_systems = set_systems;
 
         // check for conflicts
         let conflicting_systems = self.get_conflicting_systems(
@@ -1350,24 +1441,31 @@ impl ScheduleGraph {
             .zip(schedule.systems.drain(..))
             .zip(schedule.system_conditions.drain(..))
         {
-            self.systems.node_mut(key).inner = Some(system);
-            *self.systems.get_conditions_mut(key).unwrap() = conditions;
+            if let Some(node) = self.systems.node_mut(key) {
+                node.inner = Some(system);
+            }
+
+            if let Some(node_conditions) = self.systems.get_conditions_mut(key) {
+                *node_conditions = conditions;
+            }
         }
 
         for (key, conditions) in schedule
             .set_ids
             .drain(..)
             .zip(schedule.set_conditions.drain(..))
         {
-            *self.system_sets.get_conditions_mut(key).unwrap() = conditions;
+            if let Some(node_conditions) = self.system_sets.get_conditions_mut(key) {
+                *node_conditions = conditions;
+            }
         }
 
         let (new_schedule, warnings) = self.build_schedule(world, ignored_ambiguities)?;
         *schedule = new_schedule;
 
         // move systems into new schedule
         for &key in &schedule.system_ids {
-            let system = self.systems.node_mut(key).inner.take().unwrap();
+            let system = self.systems.node_mut(key).unwrap().inner.take().unwrap();
             let conditions = core::mem::take(self.systems.get_conditions_mut(key).unwrap());
             schedule.systems.push(system);
             schedule.system_conditions.push(conditions);
@@ -1766,7 +1864,7 @@ mod tests {
 
     use crate::{
         error::{ignore, panic, DefaultErrorHandler, Result},
-        prelude::{ApplyDeferred, Res, Resource},
+        prelude::{ApplyDeferred, IntoSystemSet, Res, Resource},
         schedule::{
             tests::ResMut, IntoScheduleConfigs, Schedule, ScheduleBuildSettings, SystemSet,
         },
@@ -2583,4 +2681,102 @@ mod tests {
         );
         schedule.run(&mut world);
     }
+
+    #[test]
+    fn get_a_system_key() {
+        fn test_system() {}
+
+        let mut schedule = Schedule::default();
+        schedule.add_systems(test_system);
+        let mut world = World::default();
+        let _ = schedule.initialize(&mut world);
+
+        let keys = schedule
+            .graph()
+            .systems_in_set(test_system.into_system_set().intern())
+            .unwrap();
+        assert_eq!(keys.len(), 1);
+    }
+
+    #[test]
+    fn get_system_keys_in_set() {
+        fn system_1() {}
+        fn system_2() {}
+
+        let mut schedule = Schedule::default();
+        schedule.add_systems((system_1, system_2).in_set(TestSet::First));
+        let mut world = World::default();
+        let _ = schedule.initialize(&mut world);
+
+        let keys = schedule
+            .graph()
+            .systems_in_set(TestSet::First.into_system_set().intern())
+            .unwrap();
+        assert_eq!(keys.len(), 2);
+    }
+
+    #[test]
+    fn get_system_keys_with_same_name() {
+        fn test_system() {}
+
+        let mut schedule = Schedule::default();
+        schedule.add_systems((test_system, test_system));
+        let mut world = World::default();
+        let _ = schedule.initialize(&mut world);
+
+        let keys = schedule
+            .graph()
+            .systems_in_set(test_system.into_system_set().intern())
+            .unwrap();
+        assert_ne!(keys[0], keys[1]);
+    }
+
+    #[test]
+    fn remove_a_system() {
+        fn system() {}
+
+        let mut schedule = Schedule::default();
+        schedule.add_systems(system);
+        let mut world = World::default();
+
+        let remove_count = schedule.remove_systems_in_set(system, &mut world);
+        assert_eq!(remove_count.unwrap(), 1);
+
+        // schedule has changed, so we check initializing again
+        schedule.initialize(&mut world).unwrap();
+        assert_eq!(schedule.graph().systems.len(), 0);
+    }
+
+    #[test]
+    fn remove_multiple_systems() {
+        fn system() {}
+
+        let mut schedule = Schedule::default();
+        schedule.add_systems((system, system));
+        let mut world = World::default();
+
+        let remove_count = schedule.remove_systems_in_set(system, &mut world);
+        assert_eq!(remove_count.unwrap(), 2);
+
+        // schedule has changed, so we check initializing again
+        schedule.initialize(&mut world).unwrap();
+        assert_eq!(schedule.graph().systems.len(), 0);
+    }
+
+    #[test]
+    fn remove_a_system_with_dependencies() {
+        fn system_1() {}
+        fn system_2() {}
+
+        let mut schedule = Schedule::default();
+        schedule.add_systems((system_1, system_2).chain());
+        let mut world = World::default();
+
+        let remove_count = schedule.remove_systems_in_set(system_1, &mut world);
+        assert_eq!(remove_count.unwrap(), 1);
+
+        // schedule has changed, so we check initializing again
+        schedule.initialize(&mut world).unwrap();
+        assert_eq!(schedule.graph().systems.len(), 1);
+    }
 }