From a8c610a52dd449fa4c287b4ffb9bfb83a213ebd6 Mon Sep 17 00:00:00 2001
From: Rob Grindeland <grind086@users.noreply.github.com>
Date: Tue, 12 Nov 2024 16:49:29 -0600
Subject: [PATCH] Add `unregister_system` command (#16340)

# Objective

Fixes #16266

## Solution

Added an `UnregisterSystem` command struct and
`Commands::unregister_system`. Also renamed `World::remove_system` and
`World::remove_system_cached` to `World::unregister_*`

## Testing

It's a fairly simple change, but I tested locally to ensure it actually
works.

---------

Co-authored-by: Benjamin Brienen <benjamin.brienen@outlook.com>
---
 crates/bevy_ecs/src/system/commands/mod.rs    | 12 +++++++
 crates/bevy_ecs/src/system/system_registry.rs | 36 ++++++++++++++++---
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs
index 08eda1e5d5..1480c200ab 100644
--- a/crates/bevy_ecs/src/system/commands/mod.rs
+++ b/crates/bevy_ecs/src/system/commands/mod.rs
@@ -4,6 +4,7 @@ use core::{marker::PhantomData, panic::Location};
 
 use super::{
     Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource, RunSystemCachedWith,
+    UnregisterSystem,
 };
 use crate::{
     self as bevy_ecs,
@@ -890,6 +891,17 @@ impl<'w, 's> Commands<'w, 's> {
         SystemId::from_entity(entity)
     }
 
+    /// Removes a system previously registered with [`Commands::register_system`] or [`World::register_system`].
+    ///
+    /// See [`World::unregister_system`] for more information.
+    pub fn unregister_system<I, O>(&mut self, system_id: SystemId<I, O>)
+    where
+        I: SystemInput + Send + 'static,
+        O: Send + 'static,
+    {
+        self.queue(UnregisterSystem::new(system_id));
+    }
+
     /// Similar to [`Self::run_system`], but caching the [`SystemId`] in a
     /// [`CachedSystemId`](crate::system::CachedSystemId) resource.
     ///
diff --git a/crates/bevy_ecs/src/system/system_registry.rs b/crates/bevy_ecs/src/system/system_registry.rs
index 1aa1770934..573194713f 100644
--- a/crates/bevy_ecs/src/system/system_registry.rs
+++ b/crates/bevy_ecs/src/system/system_registry.rs
@@ -30,7 +30,7 @@ pub struct SystemIdMarker;
 /// A system that has been removed from the registry.
 /// It contains the system and whether or not it has been initialized.
 ///
-/// This struct is returned by [`World::remove_system`].
+/// This struct is returned by [`World::unregister_system`].
 pub struct RemovedSystem<I = (), O = ()> {
     initialized: bool,
     system: BoxedSystem<I, O>,
@@ -172,7 +172,7 @@ impl World {
     ///
     /// If no system corresponds to the given [`SystemId`], this method returns an error.
     /// Systems are also not allowed to remove themselves, this returns an error too.
-    pub fn remove_system<I, O>(
+    pub fn unregister_system<I, O>(
         &mut self,
         id: SystemId<I, O>,
     ) -> Result<RemovedSystem<I, O>, RegisteredSystemError<I, O>>
@@ -412,7 +412,7 @@ impl World {
     /// Removes a cached system and its [`CachedSystemId`] resource.
     ///
     /// See [`World::register_system_cached`] for more information.
-    pub fn remove_system_cached<I, O, M, S>(
+    pub fn unregister_system_cached<I, O, M, S>(
         &mut self,
         _system: S,
     ) -> Result<RemovedSystem<I, O>, RegisteredSystemError<I, O>>
@@ -424,7 +424,7 @@ impl World {
         let id = self
             .remove_resource::<CachedSystemId<S::System>>()
             .ok_or(RegisteredSystemError::SystemNotCached)?;
-        self.remove_system(id.0)
+        self.unregister_system(id.0)
     }
 
     /// Runs a cached system, registering it if necessary.
@@ -544,6 +544,32 @@ where
     }
 }
 
+/// The [`Command`] type for unregistering one-shot systems from [`Commands`](crate::system::Commands).
+pub struct UnregisterSystem<I: SystemInput + 'static, O: 'static> {
+    system_id: SystemId<I, O>,
+}
+
+impl<I, O> UnregisterSystem<I, O>
+where
+    I: SystemInput + 'static,
+    O: 'static,
+{
+    /// Creates a new [`Command`] struct, which can be added to [`Commands`](crate::system::Commands).
+    pub fn new(system_id: SystemId<I, O>) -> Self {
+        Self { system_id }
+    }
+}
+
+impl<I, O> Command for UnregisterSystem<I, O>
+where
+    I: SystemInput + 'static,
+    O: 'static,
+{
+    fn apply(self, world: &mut World) {
+        let _ = world.unregister_system(self.system_id);
+    }
+}
+
 /// The [`Command`] type for running a cached one-shot system from
 /// [`Commands`](crate::system::Commands).
 ///
@@ -834,7 +860,7 @@ mod tests {
         let new = world.register_system_cached(four);
         assert_eq!(old, new);
 
-        let result = world.remove_system_cached(four);
+        let result = world.unregister_system_cached(four);
         assert!(result.is_ok());
         let new = world.register_system_cached(four);
         assert_ne!(old, new);