Add IntoSystem::with_input and ::with_input_from system wrappers (#18067)

# Objective

Originally [provided as a solution to a user's problem in
Discord](https://discord.com/channels/691052431525675048/1247654592838111302/1344431131277394042),
library authors might find the need to present user-registered systems
with system-specific data. Typically `Local<T>` is used for this type of
thing, but its not generally feasible or possible to configure/set the
underlying `T` data for locals. Alternatively, we can use `SystemInput`
to pass the data.

## Solution

- Added `IntoSystem::with_input`: Allows system-specific data to be
passed in explicitly.
- Added `IntoSystem::with_input_from`: Allows system-specific data to be
created at initialization time via `FromWorld`.

## Testing

Added two new tests, testing each of `with_input` and `with_input_from`.

Author: ItsDoot

crates/bevy_ecs/src/system/mod.rs

@@ -153,7 +153,7 @@ pub use system_name::*;
 pub use system_param::*;
 pub use system_registry::*;
 
-use crate::world::World;
+use crate::world::{FromWorld, World};
 
 /// Conversion trait to turn something into a [`System`].
 ///
@@ -228,6 +228,77 @@ pub trait IntoSystem<In: SystemInput, Out, Marker>: Sized {
         IntoAdapterSystem::new(f, self)
     }
 
+    /// Passes a mutable reference to `value` as input to the system each run,
+    /// turning it into a system that takes no input.
+    ///
+    /// `Self` can have any [`SystemInput`] type that takes a mutable reference
+    /// to `T`, such as [`InMut`].
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// fn my_system(InMut(value): InMut<usize>) {
+    ///     *value += 1;
+    ///     if *value > 10 {
+    ///        println!("Value is greater than 10!");
+    ///     }
+    /// }
+    ///
+    /// # let mut schedule = Schedule::default();
+    /// schedule.add_systems(my_system.with_input(0));
+    /// # bevy_ecs::system::assert_is_system(my_system.with_input(0));
+    /// ```
+    fn with_input<T>(self, value: T) -> WithInputWrapper<Self::System, T>
+    where
+        for<'i> In: SystemInput<Inner<'i> = &'i mut T>,
+        T: Send + Sync + 'static,
+    {
+        WithInputWrapper::new(self, value)
+    }
+
+    /// Passes a mutable reference to a value of type `T` created via
+    /// [`FromWorld`] as input to the system each run, turning it into a system
+    /// that takes no input.
+    ///
+    /// `Self` can have any [`SystemInput`] type that takes a mutable reference
+    /// to `T`, such as [`InMut`].
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// struct MyData {
+    ///     value: usize,
+    /// }
+    ///
+    /// impl FromWorld for MyData {
+    ///     fn from_world(world: &mut World) -> Self {
+    ///         // Fetch from the world the data needed to create `MyData`
+    /// #       MyData { value: 0 }
+    ///     }
+    /// }
+    ///
+    /// fn my_system(InMut(data): InMut<MyData>) {
+    ///     data.value += 1;
+    ///     if data.value > 10 {
+    ///         println!("Value is greater than 10!");
+    ///     }
+    /// }
+    /// # let mut schedule = Schedule::default();
+    /// schedule.add_systems(my_system.with_input_from::<MyData>());
+    /// # bevy_ecs::system::assert_is_system(my_system.with_input_from::<MyData>());
+    /// ```
+    fn with_input_from<T>(self) -> WithInputFromWrapper<Self::System, T>
+    where
+        for<'i> In: SystemInput<Inner<'i> = &'i mut T>,
+        T: FromWorld + Send + Sync + 'static,
+    {
+        WithInputFromWrapper::new(self)
+    }
+
     /// Get the [`TypeId`] of the [`System`] produced after calling [`into_system`](`IntoSystem::into_system`).
     #[inline]
     fn system_type_id(&self) -> TypeId {
@@ -347,8 +418,8 @@ mod tests {
             Schedule,
         },
         system::{
-            Commands, In, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query, Res, ResMut,
-            Single, StaticSystemParam, System, SystemState,
+            Commands, In, InMut, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query, Res,
+            ResMut, Single, StaticSystemParam, System, SystemState,
         },
         world::{DeferredWorld, EntityMut, FromWorld, OnAdd, World},
     };
@@ -1879,4 +1950,40 @@ mod tests {
             .commands()
             .queue(|_world: &mut World| -> () { todo!() });
     }
+
+    #[test]
+    fn with_input() {
+        fn sys(InMut(v): InMut<usize>) {
+            *v += 1;
+        }
+
+        let mut world = World::new();
+        let mut system = IntoSystem::into_system(sys.with_input(42));
+        system.initialize(&mut world);
+        system.run((), &mut world);
+        assert_eq!(*system.value(), 43);
+    }
+
+    #[test]
+    fn with_input_from() {
+        struct TestData(usize);
+
+        impl FromWorld for TestData {
+            fn from_world(_world: &mut World) -> Self {
+                Self(5)
+            }
+        }
+
+        fn sys(InMut(v): InMut<TestData>) {
+            v.0 += 1;
+        }
+
+        let mut world = World::new();
+        let mut system = IntoSystem::into_system(sys.with_input_from::<TestData>());
+        assert!(system.value().is_none());
+        system.initialize(&mut world);
+        assert!(system.value().is_some());
+        system.run((), &mut world);
+        assert_eq!(system.value().unwrap().0, 6);
+    }
 }

crates/bevy_ecs/src/system/schedule_system.rs

@@ -5,8 +5,8 @@ use crate::{
     component::{ComponentId, Tick},
     error::Result,
     query::{Access, FilteredAccessSet},
-    system::{input::SystemIn, BoxedSystem, System},
-    world::{unsafe_world_cell::UnsafeWorldCell, DeferredWorld, World},
+    system::{input::SystemIn, BoxedSystem, System, SystemInput},
+    world::{unsafe_world_cell::UnsafeWorldCell, DeferredWorld, FromWorld, World},
 };
 
 use super::{IntoSystem, SystemParamValidationError};
@@ -118,5 +118,238 @@ impl<S: System<In = ()>> System for InfallibleSystemWrapper<S> {
     }
 }
 
+/// See [`IntoSystem::with_input`] for details.
+pub struct WithInputWrapper<S, T>
+where
+    for<'i> S: System<In: SystemInput<Inner<'i> = &'i mut T>>,
+    T: Send + Sync + 'static,
+{
+    system: S,
+    value: T,
+}
+
+impl<S, T> WithInputWrapper<S, T>
+where
+    for<'i> S: System<In: SystemInput<Inner<'i> = &'i mut T>>,
+    T: Send + Sync + 'static,
+{
+    /// Wraps the given system with the given input value.
+    pub fn new<M>(system: impl IntoSystem<S::In, S::Out, M, System = S>, value: T) -> Self {
+        Self {
+            system: IntoSystem::into_system(system),
+            value,
+        }
+    }
+
+    /// Returns a reference to the input value.
+    pub fn value(&self) -> &T {
+        &self.value
+    }
+
+    /// Returns a mutable reference to the input value.
+    pub fn value_mut(&mut self) -> &mut T {
+        &mut self.value
+    }
+}
+
+impl<S, T> System for WithInputWrapper<S, T>
+where
+    for<'i> S: System<In: SystemInput<Inner<'i> = &'i mut T>>,
+    T: Send + Sync + 'static,
+{
+    type In = ();
+
+    type Out = S::Out;
+
+    fn name(&self) -> Cow<'static, str> {
+        self.system.name()
+    }
+
+    fn component_access(&self) -> &Access<ComponentId> {
+        self.system.component_access()
+    }
+
+    fn component_access_set(&self) -> &FilteredAccessSet<ComponentId> {
+        self.system.component_access_set()
+    }
+
+    fn archetype_component_access(&self) -> &Access<ArchetypeComponentId> {
+        self.system.archetype_component_access()
+    }
+
+    fn is_send(&self) -> bool {
+        self.system.is_send()
+    }
+
+    fn is_exclusive(&self) -> bool {
+        self.system.is_exclusive()
+    }
+
+    fn has_deferred(&self) -> bool {
+        self.system.has_deferred()
+    }
+
+    unsafe fn run_unsafe(
+        &mut self,
+        _input: SystemIn<'_, Self>,
+        world: UnsafeWorldCell,
+    ) -> Self::Out {
+        self.system.run_unsafe(&mut self.value, world)
+    }
+
+    fn apply_deferred(&mut self, world: &mut World) {
+        self.system.apply_deferred(world);
+    }
+
+    fn queue_deferred(&mut self, world: DeferredWorld) {
+        self.system.queue_deferred(world);
+    }
+
+    unsafe fn validate_param_unsafe(
+        &mut self,
+        world: UnsafeWorldCell,
+    ) -> Result<(), SystemParamValidationError> {
+        self.system.validate_param_unsafe(world)
+    }
+
+    fn initialize(&mut self, world: &mut World) {
+        self.system.initialize(world);
+    }
+
+    fn update_archetype_component_access(&mut self, world: UnsafeWorldCell) {
+        self.system.update_archetype_component_access(world);
+    }
+
+    fn check_change_tick(&mut self, change_tick: Tick) {
+        self.system.check_change_tick(change_tick);
+    }
+
+    fn get_last_run(&self) -> Tick {
+        self.system.get_last_run()
+    }
+
+    fn set_last_run(&mut self, last_run: Tick) {
+        self.system.set_last_run(last_run);
+    }
+}
+
+/// Constructed in [`IntoSystem::with_input_from`].
+pub struct WithInputFromWrapper<S, T> {
+    system: S,
+    value: Option<T>,
+}
+
+impl<S, T> WithInputFromWrapper<S, T>
+where
+    for<'i> S: System<In: SystemInput<Inner<'i> = &'i mut T>>,
+    T: Send + Sync + 'static,
+{
+    /// Wraps the given system.
+    pub fn new<M>(system: impl IntoSystem<S::In, S::Out, M, System = S>) -> Self {
+        Self {
+            system: IntoSystem::into_system(system),
+            value: None,
+        }
+    }
+
+    /// Returns a reference to the input value, if it has been initialized.
+    pub fn value(&self) -> Option<&T> {
+        self.value.as_ref()
+    }
+
+    /// Returns a mutable reference to the input value, if it has been initialized.
+    pub fn value_mut(&mut self) -> Option<&mut T> {
+        self.value.as_mut()
+    }
+}
+
+impl<S, T> System for WithInputFromWrapper<S, T>
+where
+    for<'i> S: System<In: SystemInput<Inner<'i> = &'i mut T>>,
+    T: FromWorld + Send + Sync + 'static,
+{
+    type In = ();
+
+    type Out = S::Out;
+
+    fn name(&self) -> Cow<'static, str> {
+        self.system.name()
+    }
+
+    fn component_access(&self) -> &Access<ComponentId> {
+        self.system.component_access()
+    }
+
+    fn component_access_set(&self) -> &FilteredAccessSet<ComponentId> {
+        self.system.component_access_set()
+    }
+
+    fn archetype_component_access(&self) -> &Access<ArchetypeComponentId> {
+        self.system.archetype_component_access()
+    }
+
+    fn is_send(&self) -> bool {
+        self.system.is_send()
+    }
+
+    fn is_exclusive(&self) -> bool {
+        self.system.is_exclusive()
+    }
+
+    fn has_deferred(&self) -> bool {
+        self.system.has_deferred()
+    }
+
+    unsafe fn run_unsafe(
+        &mut self,
+        _input: SystemIn<'_, Self>,
+        world: UnsafeWorldCell,
+    ) -> Self::Out {
+        let value = self
+            .value
+            .as_mut()
+            .expect("System input value was not found. Did you forget to initialize the system before running it?");
+        self.system.run_unsafe(value, world)
+    }
+
+    fn apply_deferred(&mut self, world: &mut World) {
+        self.system.apply_deferred(world);
+    }
+
+    fn queue_deferred(&mut self, world: DeferredWorld) {
+        self.system.queue_deferred(world);
+    }
+
+    unsafe fn validate_param_unsafe(
+        &mut self,
+        world: UnsafeWorldCell,
+    ) -> Result<(), SystemParamValidationError> {
+        self.system.validate_param_unsafe(world)
+    }
+
+    fn initialize(&mut self, world: &mut World) {
+        self.system.initialize(world);
+        if self.value.is_none() {
+            self.value = Some(T::from_world(world));
+        }
+    }
+
+    fn update_archetype_component_access(&mut self, world: UnsafeWorldCell) {
+        self.system.update_archetype_component_access(world);
+    }
+
+    fn check_change_tick(&mut self, change_tick: Tick) {
+        self.system.check_change_tick(change_tick);
+    }
+
+    fn get_last_run(&self) -> Tick {
+        self.system.get_last_run()
+    }
+
+    fn set_last_run(&mut self, last_run: Tick) {
+        self.system.set_last_run(last_run);
+    }
+}
+
 /// Type alias for a `BoxedSystem` that a `Schedule` can store.
 pub type ScheduleSystem = BoxedSystem<(), Result>;