one shot system cleanup (#16516)

# Objective

- Fixes #16497
- This is my first PR, so I'm still learning to contribute to the
project

## Solution

- Added struct `UnregisterSystemCached` and function
`unregister_system_cached`
- renamed `World::run_system_with_input` to `run_system_with`
- reordered input parameters for `World::run_system_once_with`

## Testing

- Added a crude test which registers a system via
`World::register_system_cached`, and removes it via
`Command::unregister_system_cached`.

## Migration Guide

- Change all occurrences of `World::run_system_with_input` to
`World::run_system_with`.
- swap the order of input parameters for `World::run_system_once_with`
such that the system comes before the input.

---------

Co-authored-by: Paul Mattern <[email protected]>

crates/bevy_ecs/src/system/commands/mod.rs

@@ -4,7 +4,7 @@ use core::{marker::PhantomData, panic::Location};
 
 use super::{
     Deferred, IntoObserverSystem, IntoSystem, RegisterSystem, Resource, RunSystemCachedWith,
-    UnregisterSystem,
+    UnregisterSystem, UnregisterSystemCached,
 };
 use crate::{
     self as bevy_ecs,
@@ -15,7 +15,7 @@ use crate::{
     event::{Event, SendEvent},
     observer::{Observer, TriggerEvent, TriggerTargets},
     schedule::ScheduleLabel,
-    system::{input::SystemInput, RunSystemWithInput, SystemId},
+    system::{input::SystemInput, RunSystemWith, SystemId},
     world::{
         command_queue::RawCommandQueue, unsafe_world_cell::UnsafeWorldCell, Command, CommandQueue,
         EntityWorldMut, FromWorld, SpawnBatchIter, World,
@@ -810,25 +810,25 @@ impl<'w, 's> Commands<'w, 's> {
     ///
     /// There is no way to get the output of a system when run as a command, because the
     /// execution of the system happens later. To get the output of a system, use
-    /// [`World::run_system`] or [`World::run_system_with_input`] instead of running the system as a command.
+    /// [`World::run_system`] or [`World::run_system_with`] instead of running the system as a command.
     pub fn run_system(&mut self, id: SystemId) {
-        self.run_system_with_input(id, ());
+        self.run_system_with(id, ());
     }
 
     /// Runs the system corresponding to the given [`SystemId`].
     /// Systems are ran in an exclusive and single threaded way.
     /// Running slow systems can become a bottleneck.
     ///
-    /// Calls [`World::run_system_with_input`](World::run_system_with_input).
+    /// Calls [`World::run_system_with`](World::run_system_with).
     ///
     /// There is no way to get the output of a system when run as a command, because the
     /// execution of the system happens later. To get the output of a system, use
-    /// [`World::run_system`] or [`World::run_system_with_input`] instead of running the system as a command.
-    pub fn run_system_with_input<I>(&mut self, id: SystemId<I>, input: I::Inner<'static>)
+    /// [`World::run_system`] or [`World::run_system_with`] instead of running the system as a command.
+    pub fn run_system_with<I>(&mut self, id: SystemId<I>, input: I::Inner<'static>)
     where
         I: SystemInput<Inner<'static>: Send> + 'static,
     {
-        self.queue(RunSystemWithInput::new_with_input(id, input));
+        self.queue(RunSystemWith::new_with_input(id, input));
     }
 
     /// Registers a system and returns a [`SystemId`] so it can later be called by [`World::run_system`].
@@ -904,6 +904,21 @@ impl<'w, 's> Commands<'w, 's> {
         self.queue(UnregisterSystem::new(system_id));
     }
 
+    /// Removes a system previously registered with [`World::register_system_cached`].
+    ///
+    /// See [`World::unregister_system_cached`] for more information.
+    pub fn unregister_system_cached<
+        I: SystemInput + Send + 'static,
+        O: 'static,
+        M: 'static,
+        S: IntoSystem<I, O, M> + Send + 'static,
+    >(
+        &mut self,
+        system: S,
+    ) {
+        self.queue(UnregisterSystemCached::new(system));
+    }
+
     /// Similar to [`Self::run_system`], but caching the [`SystemId`] in a
     /// [`CachedSystemId`](crate::system::CachedSystemId) resource.
     ///
@@ -915,7 +930,7 @@ impl<'w, 's> Commands<'w, 's> {
         self.run_system_cached_with(system, ());
     }
 
-    /// Similar to [`Self::run_system_with_input`], but caching the [`SystemId`] in a
+    /// Similar to [`Self::run_system_with`], but caching the [`SystemId`] in a
     /// [`CachedSystemId`](crate::system::CachedSystemId) resource.
     ///
     /// See [`World::register_system_cached`] for more information.
@@ -2617,6 +2632,25 @@ mod tests {
         assert!(world.get::<Z>(e).is_some());
     }
 
+    #[test]
+    fn unregister_system_cached_commands() {
+        let mut world = World::default();
+        let mut queue = CommandQueue::default();
+
+        fn nothing() {}
+
+        assert!(world.iter_resources().count() == 0);
+        let id = world.register_system_cached(nothing);
+        assert!(world.iter_resources().count() == 1);
+        assert!(world.get_entity(id.entity).is_ok());
+
+        let mut commands = Commands::new(&mut queue, &world);
+        commands.unregister_system_cached(nothing);
+        queue.apply(&mut world);
+        assert!(world.iter_resources().count() == 0);
+        assert!(world.get_entity(id.entity).is_err());
+    }
+
     fn is_send<T: Send>() {}
     fn is_sync<T: Sync>() {}
 

crates/bevy_ecs/src/system/system.rs

@@ -322,14 +322,14 @@ pub trait RunSystemOnce: Sized {
     where
         T: IntoSystem<(), Out, Marker>,
     {
-        self.run_system_once_with((), system)
+        self.run_system_once_with(system, ())
     }
 
     /// Tries to run a system with given input and apply deferred parameters.
     fn run_system_once_with<T, In, Out, Marker>(
         self,
-        input: SystemIn<'_, T::System>,
         system: T,
+        input: SystemIn<'_, T::System>,
     ) -> Result<Out, RunSystemError>
     where
         T: IntoSystem<In, Out, Marker>,
@@ -339,8 +339,8 @@ pub trait RunSystemOnce: Sized {
 impl RunSystemOnce for &mut World {
     fn run_system_once_with<T, In, Out, Marker>(
         self,
-        input: SystemIn<'_, T::System>,
         system: T,
+        input: SystemIn<'_, T::System>,
     ) -> Result<Out, RunSystemError>
     where
         T: IntoSystem<In, Out, Marker>,
@@ -392,7 +392,7 @@ mod tests {
         }
 
         let mut world = World::default();
-        let n = world.run_system_once_with(1, system).unwrap();
+        let n = world.run_system_once_with(system, 1).unwrap();
         assert_eq!(n, 2);
         assert_eq!(world.resource::<T>().0, 1);
     }

crates/bevy_ecs/src/system/system_registry.rs

@@ -201,7 +201,7 @@ impl World {
     /// This is different from [`RunSystemOnce::run_system_once`](crate::system::RunSystemOnce::run_system_once),
     /// because it keeps local state between calls and change detection works correctly.
     ///
-    /// In order to run a chained system with an input, use [`World::run_system_with_input`] instead.
+    /// In order to run a chained system with an input, use [`World::run_system_with`] instead.
     ///
     /// # Limitations
     ///
@@ -286,7 +286,7 @@ impl World {
         &mut self,
         id: SystemId<(), O>,
     ) -> Result<O, RegisteredSystemError<(), O>> {
-        self.run_system_with_input(id, ())
+        self.run_system_with(id, ())
     }
 
     /// Run a stored chained system by its [`SystemId`], providing an input value.
@@ -309,13 +309,13 @@ impl World {
     /// let mut world = World::default();
     /// let counter_one = world.register_system(increment);
     /// let counter_two = world.register_system(increment);
-    /// assert_eq!(world.run_system_with_input(counter_one, 1).unwrap(), 1);
-    /// assert_eq!(world.run_system_with_input(counter_one, 20).unwrap(), 21);
-    /// assert_eq!(world.run_system_with_input(counter_two, 30).unwrap(), 30);
+    /// assert_eq!(world.run_system_with(counter_one, 1).unwrap(), 1);
+    /// assert_eq!(world.run_system_with(counter_one, 20).unwrap(), 21);
+    /// assert_eq!(world.run_system_with(counter_two, 30).unwrap(), 30);
     /// ```
     ///
     /// See [`World::run_system`] for more examples.
-    pub fn run_system_with_input<I, O>(
+    pub fn run_system_with<I, O>(
         &mut self,
         id: SystemId<I, O>,
         input: I::Inner<'_>,
@@ -451,11 +451,11 @@ impl World {
         S: IntoSystem<I, O, M> + 'static,
     {
         let id = self.register_system_cached(system);
-        self.run_system_with_input(id, input)
+        self.run_system_with(id, input)
     }
 }
 
-/// The [`Command`] type for [`World::run_system`] or [`World::run_system_with_input`].
+/// The [`Command`] type for [`World::run_system`] or [`World::run_system_with`].
 ///
 /// This command runs systems in an exclusive and single threaded way.
 /// Running slow systems can become a bottleneck.
@@ -465,9 +465,9 @@ impl World {
 ///
 /// There is no way to get the output of a system when run as a command, because the
 /// execution of the system happens later. To get the output of a system, use
-/// [`World::run_system`] or [`World::run_system_with_input`] instead of running the system as a command.
+/// [`World::run_system`] or [`World::run_system_with`] instead of running the system as a command.
 #[derive(Debug, Clone)]
-pub struct RunSystemWithInput<I: SystemInput + 'static> {
+pub struct RunSystemWith<I: SystemInput + 'static> {
     system_id: SystemId<I>,
     input: I::Inner<'static>,
 }
@@ -478,12 +478,12 @@ pub struct RunSystemWithInput<I: SystemInput + 'static> {
 /// Running slow systems can become a bottleneck.
 ///
 /// If the system needs an [`In<_>`](crate::system::In) input value to run, use the
-/// [`RunSystemWithInput`] type instead.
+/// [`RunSystemWith`] type instead.
 ///
 /// There is no way to get the output of a system when run as a command, because the
 /// execution of the system happens later. To get the output of a system, use
-/// [`World::run_system`] or [`World::run_system_with_input`] instead of running the system as a command.
-pub type RunSystem = RunSystemWithInput<()>;
+/// [`World::run_system`] or [`World::run_system_with`] instead of running the system as a command.
+pub type RunSystem = RunSystemWith<()>;
 
 impl RunSystem {
     /// Creates a new [`Command`] struct, which can be added to [`Commands`](crate::system::Commands).
@@ -492,21 +492,21 @@ impl RunSystem {
     }
 }
 
-impl<I: SystemInput + 'static> RunSystemWithInput<I> {
+impl<I: SystemInput + 'static> RunSystemWith<I> {
     /// Creates a new [`Command`] struct, which can be added to [`Commands`](crate::system::Commands)
     /// in order to run the specified system with the provided [`In<_>`](crate::system::In) input value.
     pub fn new_with_input(system_id: SystemId<I>, input: I::Inner<'static>) -> Self {
         Self { system_id, input }
     }
 }
 
-impl<I> Command for RunSystemWithInput<I>
+impl<I> Command for RunSystemWith<I>
 where
     I: SystemInput<Inner<'static>: Send> + 'static,
 {
     #[inline]
     fn apply(self, world: &mut World) {
-        _ = world.run_system_with_input(self.system_id, self.input);
+        _ = world.run_system_with(self.system_id, self.input);
     }
 }
 
@@ -570,6 +570,42 @@ where
     }
 }
 
+/// The [`Command`] type for unregistering one-shot systems from [`Commands`](crate::system::Commands).
+pub struct UnregisterSystemCached<I, O, M, S>
+where
+    I: SystemInput + 'static,
+    S: IntoSystem<I, O, M> + Send + 'static,
+{
+    system: S,
+    _phantom: PhantomData<fn() -> (I, O, M)>,
+}
+
+impl<I, O, M, S> UnregisterSystemCached<I, O, M, S>
+where
+    I: SystemInput + 'static,
+    S: IntoSystem<I, O, M> + Send + 'static,
+{
+    /// Creates a new [`Command`] struct, which can be added to [`Commands`](crate::system::Commands).
+    pub fn new(system: S) -> Self {
+        Self {
+            system,
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<I, O, M, S> Command for UnregisterSystemCached<I, O, M, S>
+where
+    I: SystemInput + 'static,
+    O: 'static,
+    M: 'static,
+    S: IntoSystem<I, O, M> + Send + 'static,
+{
+    fn apply(self, world: &mut World) {
+        let _ = world.unregister_system_cached(self.system);
+    }
+}
+
 /// The [`Command`] type for running a cached one-shot system from
 /// [`Commands`](crate::system::Commands).
 ///
@@ -730,22 +766,22 @@ mod tests {
         assert_eq!(*world.resource::<Counter>(), Counter(1));
 
         world
-            .run_system_with_input(id, NonCopy(1))
+            .run_system_with(id, NonCopy(1))
             .expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(2));
 
         world
-            .run_system_with_input(id, NonCopy(1))
+            .run_system_with(id, NonCopy(1))
             .expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(3));
 
         world
-            .run_system_with_input(id, NonCopy(20))
+            .run_system_with(id, NonCopy(20))
             .expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(23));
 
         world
-            .run_system_with_input(id, NonCopy(1))
+            .run_system_with(id, NonCopy(1))
             .expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(24));
     }
@@ -828,7 +864,7 @@ mod tests {
 
         fn nested(query: Query<&Callback>, mut commands: Commands) {
             for callback in query.iter() {
-                commands.run_system_with_input(callback.0, callback.1);
+                commands.run_system_with(callback.0, callback.1);
             }
         }
 
@@ -922,7 +958,7 @@ mod tests {
         world.insert_resource(Counter(0));
 
         let id = world.register_system(with_ref);
-        world.run_system_with_input(id, &2).unwrap();
+        world.run_system_with(id, &2).unwrap();
         assert_eq!(*world.resource::<Counter>(), Counter(2));
     }
 
@@ -944,15 +980,11 @@ mod tests {
         let post_system = world.register_system(post);
 
         let mut event = MyEvent { cancelled: false };
-        world
-            .run_system_with_input(post_system, &mut event)
-            .unwrap();
+        world.run_system_with(post_system, &mut event).unwrap();
         assert!(!event.cancelled);
 
         world.resource_mut::<Counter>().0 = 1;
-        world
-            .run_system_with_input(post_system, &mut event)
-            .unwrap();
+        world.run_system_with(post_system, &mut event).unwrap();
         assert!(event.cancelled);
     }
 

crates/bevy_remote/src/lib.rs

@@ -800,7 +800,7 @@ fn process_remote_requests(world: &mut World) {
 
         match handler {
             RemoteMethodSystemId::Instant(id) => {
-                let result = match world.run_system_with_input(id, message.params) {
+                let result = match world.run_system_with(id, message.params) {
                     Ok(result) => result,
                     Err(error) => {
                         let _ = message.sender.force_send(Err(BrpError {
@@ -850,7 +850,7 @@ fn process_single_ongoing_watching_request(
     system_id: &RemoteWatchingMethodSystemId,
 ) -> BrpResult<Option<Value>> {
     world
-        .run_system_with_input(*system_id, message.params.clone())
+        .run_system_with(*system_id, message.params.clone())
         .map_err(|error| BrpError {
             code: error_codes::INTERNAL_ERROR,
             message: format!("Failed to run method handler: {error}"),

crates/bevy_ui/src/layout/mod.rs

@@ -944,7 +944,7 @@ mod tests {
             new_pos: Vec2,
             expected_camera_entity: &Entity,
         ) {
-            world.run_system_once_with(new_pos, move_ui_node).unwrap();
+            world.run_system_once_with(move_ui_node, new_pos).unwrap();
             ui_schedule.run(world);
             let (ui_node_entity, TargetCamera(target_camera_entity)) = world
                 .query_filtered::<(Entity, &TargetCamera), With<MovingUiNode>>()
@@ -1234,11 +1234,11 @@ mod tests {
         }
 
         let _ = world.run_system_once_with(
+            test_system,
             TestSystemParam {
                 camera_entity,
                 root_node_entity,
             },
-            test_system,
         );
 
         let ui_surface = world.resource::<UiSurface>();