From 90eaaa93b29614a2e76f5a542ec7429ec79b3e0b Mon Sep 17 00:00:00 2001
From: Nathan Fenner <nfenneremail@gmail.com>
Date: Sun, 12 Nov 2023 15:04:23 -0800
Subject: [PATCH] support the run_system_with_input Command

---
 crates/bevy_ecs/src/system/commands/mod.rs    | 14 +++-
 crates/bevy_ecs/src/system/system_registry.rs | 69 +++++++++++++++++--
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/crates/bevy_ecs/src/system/commands/mod.rs b/crates/bevy_ecs/src/system/commands/mod.rs
index 89c8493dc667b..304ccf3bbe2e0 100644
--- a/crates/bevy_ecs/src/system/commands/mod.rs
+++ b/crates/bevy_ecs/src/system/commands/mod.rs
@@ -5,7 +5,7 @@ use crate::{
     self as bevy_ecs,
     bundle::Bundle,
     entity::{Entities, Entity},
-    system::{RunSystem, SystemId},
+    system::{RunSystemWithInput, SystemId},
     world::{EntityWorldMut, FromWorld, World},
 };
 use bevy_ecs_macros::SystemParam;
@@ -524,7 +524,17 @@ impl<'w, 's> Commands<'w, 's> {
     ///
     /// Calls [`World::run_system`](crate::system::World::run_system).
     pub fn run_system(&mut self, id: SystemId) {
-        self.queue.push(RunSystem::new(id));
+        self.run_system_with_input(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`](crate::system::World::run_system_with_input).
+    pub fn run_system_with_input<I: 'static + Send>(&mut self, id: SystemId<I>, input: I) {
+        self.queue
+            .push(RunSystemWithInput::new_with_input(id, input));
     }
 
     /// Pushes a generic [`Command`] to the command queue.
diff --git a/crates/bevy_ecs/src/system/system_registry.rs b/crates/bevy_ecs/src/system/system_registry.rs
index 99262271e9e17..92a8708f138ce 100644
--- a/crates/bevy_ecs/src/system/system_registry.rs
+++ b/crates/bevy_ecs/src/system/system_registry.rs
@@ -299,26 +299,55 @@ impl World {
     }
 }
 
-/// The [`Command`] type for [`World::run_system`].
+/// The [`Command`] type for [`World::run_system`] or [`World::run_system_with_input`].
 ///
 /// This command runs systems in an exclusive and single threaded way.
 /// Running slow systems can become a bottleneck.
+///
+/// If the system needs an [`In<_>`](crate::system::In) input value to run, it must
+/// be provided as part of the command.
+///
+/// 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.
 #[derive(Debug, Clone)]
-pub struct RunSystem {
-    system_id: SystemId,
+pub struct RunSystemWithInput<I: 'static> {
+    system_id: SystemId<I>,
+    input: I,
 }
 
+/// The [`Command`] type for [`World::run_system`].
+///
+/// This command runs systems in an exclusive and single threaded way.
+/// Running slow systems can become a bottleneck.
+///
+/// If the system needs an [`In<_>`](crate::system::In) input value to run, use the
+/// [crate::system::RunSystemWithInput] 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<()>;
+
 impl RunSystem {
     /// Creates a new [`Command`] struct, which can be added to [`Commands`](crate::system::Commands)
     pub fn new(system_id: SystemId) -> Self {
-        Self { system_id }
+        Self::new_with_input(system_id, ())
+    }
+}
+
+impl<I: 'static> RunSystemWithInput<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) -> Self {
+        Self { system_id, input }
     }
 }
 
-impl Command for RunSystem {
+impl<I: 'static + Send> Command for RunSystemWithInput<I> {
     #[inline]
     fn apply(self, world: &mut World) {
-        let _ = world.run_system(self.system_id);
+        let _ = world.run_system_with_input(self.system_id, self.input);
     }
 }
 
@@ -505,4 +534,32 @@ mod tests {
         let _ = world.run_system(nested_id);
         assert_eq!(*world.resource::<Counter>(), Counter(5));
     }
+
+    #[test]
+    fn nested_systems_with_inputs() {
+        use crate::system::SystemId;
+
+        #[derive(Component)]
+        struct Callback(SystemId<u8>, u8);
+
+        fn nested(query: Query<&Callback>, mut commands: Commands) {
+            for callback in query.iter() {
+                commands.run_system_with_input(callback.0, callback.1);
+            }
+        }
+
+        let mut world = World::new();
+        world.insert_resource(Counter(0));
+
+        let increment_by =
+            world.register_system(|In(amt): In<u8>, mut counter: ResMut<Counter>| {
+                counter.0 += amt;
+            });
+        let nested_id = world.register_system(nested);
+
+        world.spawn(Callback(increment_by, 2));
+        world.spawn(Callback(increment_by, 3));
+        let _ = world.run_system(nested_id);
+        assert_eq!(*world.resource::<Counter>(), Counter(5));
+    }
 }