Remove systems from schedule (#20298)

# Objective

- Add ability to remove systems from a schedule.
- Fixes #279.

## Solution

- add `remove_systems_in_set` method which removes all the systems by
their set name. In most cases systems are added to a schedule once, so
just passing the system (which is an implicit set by it's type name)
will just remove the one system. If a system is added multiple times
such as `apply_deferred`, then you may need to give the system a unique
set, if you just want to remove just one system.
- The method also tries to remove all the configuration for the set and
the systems found.
- We need to pass world to the method because user build passes are
allowed to write to the world.

## Testing

- Added some unit tests.

## Future work

- In the future of systems as entities this will probably look much
different. We'd probably store all the edges as relationships, so
hopefully cleanup would be easier.
- It'd be better if remove_systems_in_set could take a tuple of system
sets. I looked into using IntoScheduleConfigs, but that won't work since
the system set version of it errors if you try to pass a system type
set. This would allow for less rebuilding of the schedule. The schedule
needs to be rebuilt to each call to be able to find the system set.

---------

Co-authored-by: Barrett Ray <66580279+onkoe@users.noreply.github.com>
Co-authored-by: Chris Russell <8494645+chescock@users.noreply.github.com>

Author: 3 people

crates/bevy_app/src/app.rs

@@ -15,7 +15,10 @@ use bevy_ecs::{
     intern::Interned,
     message::{message_update_system, MessageCursor},
     prelude::*,
-    schedule::{InternedSystemSet, ScheduleBuildSettings, ScheduleLabel},
+    schedule::{
+        InternedSystemSet, ScheduleBuildSettings, ScheduleCleanupPolicy, ScheduleError,
+        ScheduleLabel,
+    },
     system::{IntoObserverSystem, ScheduleSystem, SystemId, SystemInput},
 };
 use bevy_platform::collections::HashMap;
@@ -325,6 +328,38 @@ impl App {
         self
     }
 
+    /// Removes all systems in a [`SystemSet`]. This will cause the schedule to be rebuilt when
+    /// the schedule is run again and can be slow. 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_app::prelude::*;
+    /// # use bevy_ecs::schedule::ScheduleCleanupPolicy;
+    /// #
+    /// # let mut app = App::new();
+    /// # fn system_a() {}
+    /// # fn system_b() {}
+    /// #
+    /// // add the system
+    /// app.add_systems(Update, system_a);
+    ///
+    /// // remove the system
+    /// app.remove_systems_in_set(Update, system_a, ScheduleCleanupPolicy::RemoveSystemsOnly);
+    /// ```
+    pub fn remove_systems_in_set<M>(
+        &mut self,
+        schedule: impl ScheduleLabel,
+        set: impl IntoSystemSet<M>,
+        policy: ScheduleCleanupPolicy,
+    ) -> Result<usize, ScheduleError> {
+        self.main_mut().remove_systems_in_set(schedule, set, policy)
+    }
+
     /// Registers a system and returns a [`SystemId`] so it can later be called by [`World::run_system`].
     ///
     /// It's possible to register the same systems more than once, they'll be stored separately.

crates/bevy_app/src/sub_app.rs

@@ -3,7 +3,10 @@ use alloc::{boxed::Box, string::String, vec::Vec};
 use bevy_ecs::{
     message::MessageRegistry,
     prelude::*,
-    schedule::{InternedScheduleLabel, InternedSystemSet, ScheduleBuildSettings, ScheduleLabel},
+    schedule::{
+        InternedScheduleLabel, InternedSystemSet, ScheduleBuildSettings, ScheduleCleanupPolicy,
+        ScheduleError, ScheduleLabel,
+    },
     system::{ScheduleSystem, SystemId, SystemInput},
 };
 use bevy_platform::collections::{HashMap, HashSet};
@@ -219,6 +222,18 @@ impl SubApp {
         self
     }
 
+    /// See [`App::remove_systems_in_set`]
+    pub fn remove_systems_in_set<M>(
+        &mut self,
+        schedule: impl ScheduleLabel,
+        set: impl IntoSystemSet<M>,
+        policy: ScheduleCleanupPolicy,
+    ) -> Result<usize, ScheduleError> {
+        self.world.schedule_scope(schedule, |world, schedule| {
+            schedule.remove_systems_in_set(set, world, policy)
+        })
+    }
+
     /// See [`App::register_system`].
     pub fn register_system<I, O, M>(
         &mut self,

crates/bevy_ecs/src/schedule/error.rs

@@ -259,3 +259,26 @@ 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 set
+    #[error("Set not found")]
+    SetNotFound,
+    /// Schedule not found
+    #[error("Schedule not found.")]
+    ScheduleNotFound,
+    /// Error initializing schedule
+    #[error("{0}")]
+    ScheduleBuildError(ScheduleBuildError),
+}
+
+impl From<ScheduleBuildError> for ScheduleError {
+    fn from(value: ScheduleBuildError) -> Self {
+        Self::ScheduleBuildError(value)
+    }
+}

crates/bevy_ecs/src/schedule/node.rs

@@ -487,14 +487,10 @@ impl Systems {
         self.nodes.get_mut(key).and_then(|node| node.get_mut())
     }
 
-    /// Returns a mutable reference to the system with the given key, panicking
-    /// if it does not exist.
-    ///
-    /// # 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]
+    /// Returns a mutable reference to the system with the given key. Will return
+    /// `None` if the key does not exist.
+    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 +550,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 +661,11 @@ impl SystemSets {
         self.sets.get(key).map(|set| &**set)
     }
 
+    /// Returns the key 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 +736,14 @@ 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);
+        self.uninit.retain(|uninit| uninit.key != key);
+        true
+    }
+
     /// Returns `true` if all system sets' conditions in this container have
     /// been initialized.
     pub fn is_initialized(&self) -> bool {