Add 'World::run_system_with_input' function + allow World::run_system to get system output (#10380)

# Objective

Allows chained systems taking an `In<_>` input parameter to be run as
one-shot systems. This API was mentioned in #8963.

In addition, `run_system(_with_input)` returns the system output, for
any `'static` output type.

## Solution

A new function, `World::run_system_with_input` allows a `SystemId<I, O>`
to be run by providing an `I` value as input and producing `O` as an
output.

`SystemId<I, O>` is now generic over the input type `I` and output type
`O`, along with the related functions and types `RegisteredSystem`,
`RemovedSystem`, `register_system`, `remove_system`, and
`RegisteredSystemError`. These default to `()`, preserving the existing
API, for all of the public types.

---

## Changelog

- Added `World::run_system_with_input` function to allow one-shot
systems that take `In<_>` input parameters
- Changed `World::run_system` and `World::register_system` to support
systems with return types beyond `()`
- Added `Commands::run_system_with_input` command that schedules a
one-shot system with an `In<_>` input parameter

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;
@@ -523,8 +523,26 @@ impl<'w, 's> Commands<'w, 's> {
     /// Running slow systems can become a bottleneck.
     ///
     /// Calls [`World::run_system`](crate::system::World::run_system).
+    ///
+    /// 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(&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).
+    ///
+    /// 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: '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.

crates/bevy_ecs/src/system/system_registry.rs

@@ -7,18 +7,18 @@ use thiserror::Error;
 
 /// A small wrapper for [`BoxedSystem`] that also keeps track whether or not the system has been initialized.
 #[derive(Component)]
-struct RegisteredSystem {
+struct RegisteredSystem<I, O> {
     initialized: bool,
-    system: BoxedSystem,
+    system: BoxedSystem<I, O>,
 }
 
 /// 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`].
-pub struct RemovedSystem {
+pub struct RemovedSystem<I = (), O = ()> {
     initialized: bool,
-    system: BoxedSystem,
+    system: BoxedSystem<I, O>,
 }
 
 impl RemovedSystem {
@@ -38,8 +38,41 @@ impl RemovedSystem {
 ///
 /// These are opaque identifiers, keyed to a specific [`World`],
 /// and are created via [`World::register_system`].
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Eq)]
-pub struct SystemId(Entity);
+pub struct SystemId<I = (), O = ()>(Entity, std::marker::PhantomData<fn(I) -> O>);
+
+// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters.
+impl<I, O> Copy for SystemId<I, O> {}
+
+// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters.
+impl<I, O> Clone for SystemId<I, O> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters.
+impl<I, O> PartialEq for SystemId<I, O> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0 == other.0 && self.1 == other.1
+    }
+}
+
+// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters.
+impl<I, O> std::hash::Hash for SystemId<I, O> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.0.hash(state);
+    }
+}
+
+impl<I, O> std::fmt::Debug for SystemId<I, O> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("SystemId")
+            .field(&self.0)
+            .field(&self.1)
+            .finish()
+    }
+}
 
 impl World {
     /// Registers a system and returns a [`SystemId`] so it can later be called by [`World::run_system`].
@@ -51,10 +84,10 @@ impl World {
     /// This allows for running systems in a pushed-based fashion.
     /// Using a [`Schedule`](crate::schedule::Schedule) is still preferred for most cases
     /// due to its better performance and abillity to run non-conflicting systems simultaneously.
-    pub fn register_system<M, S: IntoSystem<(), (), M> + 'static>(
+    pub fn register_system<I: 'static, O: 'static, M, S: IntoSystem<I, O, M> + 'static>(
         &mut self,
         system: S,
-    ) -> SystemId {
+    ) -> SystemId<I, O> {
         self.register_boxed_system(Box::new(IntoSystem::into_system(system)))
     }
 
@@ -62,13 +95,17 @@ impl World {
     ///
     ///  This is useful if the [`IntoSystem`] implementor has already been turned into a
     /// [`System`](crate::system::System) trait object and put in a [`Box`].
-    pub fn register_boxed_system(&mut self, system: BoxedSystem) -> SystemId {
+    pub fn register_boxed_system<I: 'static, O: 'static>(
+        &mut self,
+        system: BoxedSystem<I, O>,
+    ) -> SystemId<I, O> {
         SystemId(
             self.spawn(RegisteredSystem {
                 initialized: false,
                 system,
             })
             .id(),
+            std::marker::PhantomData,
         )
     }
 
@@ -78,11 +115,14 @@ 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(&mut self, id: SystemId) -> Result<RemovedSystem, RegisteredSystemError> {
+    pub fn remove_system<I: 'static, O: 'static>(
+        &mut self,
+        id: SystemId<I, O>,
+    ) -> Result<RemovedSystem<I, O>, RegisteredSystemError<I, O>> {
         match self.get_entity_mut(id.0) {
             Some(mut entity) => {
                 let registered_system = entity
-                    .take::<RegisteredSystem>()
+                    .take::<RegisteredSystem<I, O>>()
                     .ok_or(RegisteredSystemError::SelfRemove(id))?;
                 entity.despawn();
                 Ok(RemovedSystem {
@@ -100,14 +140,17 @@ 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.
+    ///
     /// # Limitations
     ///
-    ///  - Stored systems cannot be chained: they can neither have an [`In`](crate::system::In) nor return any values.
     ///  - Stored systems cannot be recursive, they cannot call themselves through [`Commands::run_system`](crate::system::Commands).
     ///  - Exclusive systems cannot be used.
     ///
     /// # Examples
     ///
+    /// ## Running a system
+    ///
     /// ```rust
     /// # use bevy_ecs::prelude::*;
     /// #[derive(Resource, Default)]
@@ -126,7 +169,7 @@ impl World {
     /// world.run_system(counter_two); // -> 1
     /// ```
     ///
-    /// Change detection:
+    /// ## Change detection
     ///
     /// ```rust
     /// # use bevy_ecs::prelude::*;
@@ -149,7 +192,81 @@ impl World {
     /// world.resource_mut::<ChangeDetector>().set_changed();
     /// let _ = world.run_system(detector); // -> Something happened!
     /// ```
-    pub fn run_system(&mut self, id: SystemId) -> Result<(), RegisteredSystemError> {
+    ///
+    /// ## Getting system output
+    ///
+    /// ```rust
+    /// # use bevy_ecs::prelude::*;
+    ///
+    /// #[derive(Resource)]
+    /// struct PlayerScore(i32);
+    ///
+    /// #[derive(Resource)]
+    /// struct OpponentScore(i32);
+    ///
+    /// fn get_player_score(player_score: Res<PlayerScore>) -> i32 {
+    ///   player_score.0
+    /// }
+    ///
+    /// fn get_opponent_score(opponent_score: Res<OpponentScore>) -> i32 {
+    ///   opponent_score.0
+    /// }
+    ///
+    /// let mut world = World::default();
+    /// world.insert_resource(PlayerScore(3));
+    /// world.insert_resource(OpponentScore(2));
+    ///
+    /// let scoring_systems = [
+    ///   ("player", world.register_system(get_player_score)),
+    ///   ("opponent", world.register_system(get_opponent_score)),
+    /// ];
+    ///
+    /// for (label, scoring_system) in scoring_systems {
+    ///   println!("{label} has score {}", world.run_system(scoring_system).expect("system succeeded"));
+    /// }
+    /// ```
+    pub fn run_system<O: 'static>(
+        &mut self,
+        id: SystemId<(), O>,
+    ) -> Result<O, RegisteredSystemError<(), O>> {
+        self.run_system_with_input(id, ())
+    }
+
+    /// Run a stored chained system by its [`SystemId`], providing an input value.
+    /// Before running a system, it must first be registered.
+    /// The method [`World::register_system`] stores a given system and returns a [`SystemId`].
+    ///
+    /// # Limitations
+    ///
+    ///  - Stored systems cannot be recursive, they cannot call themselves through [`Commands::run_system`](crate::system::Commands).
+    ///  - Exclusive systems cannot be used.
+    ///
+    /// # Examples
+    ///
+    /// ```rust
+    /// # use bevy_ecs::prelude::*;
+    /// #[derive(Resource, Default)]
+    /// struct Counter(u8);
+    ///
+    /// fn increment(In(increment_by): In<u8>, mut counter: Local<Counter>) {
+    ///    counter.0 += increment_by;
+    ///    println!("{}", counter.0);
+    /// }
+    ///
+    /// let mut world = World::default();
+    /// let counter_one = world.register_system(increment);
+    /// let counter_two = world.register_system(increment);
+    /// world.run_system_with_input(counter_one, 1); // -> 1
+    /// world.run_system_with_input(counter_one, 20); // -> 21
+    /// world.run_system_with_input(counter_two, 30); // -> 51
+    /// ```
+    ///
+    /// See [`World::run_system`] for more examples.
+    pub fn run_system_with_input<I: 'static, O: 'static>(
+        &mut self,
+        id: SystemId<I, O>,
+        input: I,
+    ) -> Result<O, RegisteredSystemError<I, O>> {
         // lookup
         let mut entity = self
             .get_entity_mut(id.0)
@@ -160,7 +277,7 @@ impl World {
             mut initialized,
             mut system,
         } = entity
-            .take::<RegisteredSystem>()
+            .take::<RegisteredSystem<I, O>>()
             .ok_or(RegisteredSystemError::Recursive(id))?;
 
         // run the system
@@ -168,57 +285,98 @@ impl World {
             system.initialize(self);
             initialized = true;
         }
-        system.run((), self);
+        let result = system.run(input, self);
         system.apply_deferred(self);
 
         // return ownership of system trait object (if entity still exists)
         if let Some(mut entity) = self.get_entity_mut(id.0) {
-            entity.insert::<RegisteredSystem>(RegisteredSystem {
+            entity.insert::<RegisteredSystem<I, O>>(RegisteredSystem {
                 initialized,
                 system,
             });
         }
-        Ok(())
+        Ok(result)
     }
 }
 
+/// 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 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.
-#[derive(Debug, Clone)]
+///
-pub struct RunSystem {
+/// If the system needs an [`In<_>`](crate::system::In) input value to run, use the
-    system_id: SystemId,
+/// [`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 Command for RunSystem {
+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<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);
     }
 }
 
 /// An operation with stored systems failed.
-#[derive(Debug, Error)]
+#[derive(Error)]
-pub enum RegisteredSystemError {
+pub enum RegisteredSystemError<I = (), O = ()> {
     /// A system was run by id, but no system with that id was found.
     ///
     /// Did you forget to register it?
     #[error("System {0:?} was not registered")]
-    SystemIdNotRegistered(SystemId),
+    SystemIdNotRegistered(SystemId<I, O>),
     /// A system tried to run itself recursively.
     #[error("System {0:?} tried to run itself recursively")]
-    Recursive(SystemId),
+    Recursive(SystemId<I, O>),
     /// A system tried to remove itself.
     #[error("System {0:?} tried to remove itself")]
-    SelfRemove(SystemId),
+    SelfRemove(SystemId<I, O>),
+}
+
+impl<I, O> std::fmt::Debug for RegisteredSystemError<I, O> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::SystemIdNotRegistered(arg0) => {
+                f.debug_tuple("SystemIdNotRegistered").field(arg0).finish()
+            }
+            Self::Recursive(arg0) => f.debug_tuple("Recursive").field(arg0).finish(),
+            Self::SelfRemove(arg0) => f.debug_tuple("SelfRemove").field(arg0).finish(),
+        }
+    }
 }
 
 mod tests {
@@ -248,14 +406,14 @@ mod tests {
         assert_eq!(*world.resource::<Counter>(), Counter(0));
         // Resources are changed when they are first added.
         let id = world.register_system(count_up_iff_changed);
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(1));
         // Nothing changed
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(1));
         // Making a change
         world.resource_mut::<ChangeDetector>().set_changed();
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(2));
     }
 
@@ -271,16 +429,82 @@ mod tests {
         world.insert_resource(Counter(1));
         assert_eq!(*world.resource::<Counter>(), Counter(1));
         let id = world.register_system(doubling);
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(1));
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(2));
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(4));
-        let _ = world.run_system(id);
+        world.run_system(id).expect("system runs successfully");
         assert_eq!(*world.resource::<Counter>(), Counter(8));
     }
 
+    #[test]
+    fn input_values() {
+        // Verify that a non-Copy, non-Clone type can be passed in.
+        struct NonCopy(u8);
+
+        fn increment_sys(In(NonCopy(increment_by)): In<NonCopy>, mut counter: ResMut<Counter>) {
+            counter.0 += increment_by;
+        }
+
+        let mut world = World::new();
+
+        let id = world.register_system(increment_sys);
+
+        // Insert the resource after registering the system.
+        world.insert_resource(Counter(1));
+        assert_eq!(*world.resource::<Counter>(), Counter(1));
+
+        world
+            .run_system_with_input(id, NonCopy(1))
+            .expect("system runs successfully");
+        assert_eq!(*world.resource::<Counter>(), Counter(2));
+
+        world
+            .run_system_with_input(id, NonCopy(1))
+            .expect("system runs successfully");
+        assert_eq!(*world.resource::<Counter>(), Counter(3));
+
+        world
+            .run_system_with_input(id, NonCopy(20))
+            .expect("system runs successfully");
+        assert_eq!(*world.resource::<Counter>(), Counter(23));
+
+        world
+            .run_system_with_input(id, NonCopy(1))
+            .expect("system runs successfully");
+        assert_eq!(*world.resource::<Counter>(), Counter(24));
+    }
+
+    #[test]
+    fn output_values() {
+        // Verify that a non-Copy, non-Clone type can be returned.
+        #[derive(Eq, PartialEq, Debug)]
+        struct NonCopy(u8);
+
+        fn increment_sys(mut counter: ResMut<Counter>) -> NonCopy {
+            counter.0 += 1;
+            NonCopy(counter.0)
+        }
+
+        let mut world = World::new();
+
+        let id = world.register_system(increment_sys);
+
+        // Insert the resource after registering the system.
+        world.insert_resource(Counter(1));
+        assert_eq!(*world.resource::<Counter>(), Counter(1));
+
+        let output = world.run_system(id).expect("system runs successfully");
+        assert_eq!(*world.resource::<Counter>(), Counter(2));
+        assert_eq!(output, NonCopy(2));
+
+        let output = world.run_system(id).expect("system runs successfully");
+        assert_eq!(*world.resource::<Counter>(), Counter(3));
+        assert_eq!(output, NonCopy(3));
+    }
+
     #[test]
     fn nested_systems() {
         use crate::system::SystemId;
@@ -310,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));
+    }
 }