diff --git a/crates/bevy_ecs/src/system/mod.rs b/crates/bevy_ecs/src/system/mod.rs index 1bdd26add2..7ba3dba4fb 100644 --- a/crates/bevy_ecs/src/system/mod.rs +++ b/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: 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) { + /// *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(self, value: T) -> WithInputWrapper + where + for<'i> In: SystemInput = &'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) { + /// 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::()); + /// # bevy_ecs::system::assert_is_system(my_system.with_input_from::()); + /// ``` + fn with_input_from(self) -> WithInputFromWrapper + where + for<'i> In: SystemInput = &'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) { + *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) { + v.0 += 1; + } + + let mut world = World::new(); + let mut system = IntoSystem::into_system(sys.with_input_from::()); + 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); + } } diff --git a/crates/bevy_ecs/src/system/schedule_system.rs b/crates/bevy_ecs/src/system/schedule_system.rs index 4ad990b47a..8ef7d9ed57 100644 --- a/crates/bevy_ecs/src/system/schedule_system.rs +++ b/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> System for InfallibleSystemWrapper { } } +/// See [`IntoSystem::with_input`] for details. +pub struct WithInputWrapper +where + for<'i> S: System = &'i mut T>>, + T: Send + Sync + 'static, +{ + system: S, + value: T, +} + +impl WithInputWrapper +where + for<'i> S: System = &'i mut T>>, + T: Send + Sync + 'static, +{ + /// Wraps the given system with the given input value. + pub fn new(system: impl IntoSystem, 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 System for WithInputWrapper +where + for<'i> S: System = &'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 { + self.system.component_access() + } + + fn component_access_set(&self) -> &FilteredAccessSet { + self.system.component_access_set() + } + + fn archetype_component_access(&self) -> &Access { + 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 { + system: S, + value: Option, +} + +impl WithInputFromWrapper +where + for<'i> S: System = &'i mut T>>, + T: Send + Sync + 'static, +{ + /// Wraps the given system. + pub fn new(system: impl IntoSystem) -> 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 System for WithInputFromWrapper +where + for<'i> S: System = &'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 { + self.system.component_access() + } + + fn component_access_set(&self) -> &FilteredAccessSet { + self.system.component_access_set() + } + + fn archetype_component_access(&self) -> &Access { + 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>;