From b56de8d788552b69c4f3950beb986337db3c87cc Mon Sep 17 00:00:00 2001 From: Chris Russell <8494645+chescock@users.noreply.github.com> Date: Fri, 7 Jun 2024 12:09:03 -0400 Subject: [PATCH] Support dynamic system parameters. --- crates/bevy_ecs/src/system/builder.rs | 62 +++++++- crates/bevy_ecs/src/system/system_param.rs | 166 +++++++++++++++++++++ 2 files changed, 227 insertions(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/system/builder.rs b/crates/bevy_ecs/src/system/builder.rs index bb4584b31c32f..15f08f197b534 100644 --- a/crates/bevy_ecs/src/system/builder.rs +++ b/crates/bevy_ecs/src/system/builder.rs @@ -4,7 +4,7 @@ use crate::{ prelude::QueryBuilder, query::{QueryData, QueryFilter, QueryState}, system::{ - system_param::{Local, ParamSet, SystemParam}, + system_param::{DynSystemParam, DynSystemParamState, Local, ParamSet, SystemParam}, Query, SystemMeta, }, world::{FromWorld, World}, @@ -292,6 +292,34 @@ unsafe impl<'w, 's, P: SystemParam, B: SystemParamBuilder

> } } +/// A [`SystemParamBuilder`] for a [`DynSystemParam`]. +pub struct DynParamBuilder<'a>( + Box DynSystemParamState + 'a>, +); + +impl<'a> DynParamBuilder<'a> { + /// Creates a new [`DynParamBuilder`] by wrapping a [`SystemParamBuilder`] of any type. + /// The built [`DynSystemParam`] can be downcast to `T`. + pub fn new(builder: impl SystemParamBuilder + 'a) -> Self { + Self(Box::new(|world, meta| { + DynSystemParamState::new::(builder.build(world, meta)) + })) + } +} + +// SAFETY: `DynSystemParam::get_param` will call `get_param` on the boxed `DynSystemParamState`, +// and the boxed builder was a valid implementation of `SystemParamBuilder` for that type. +// The resulting `DynSystemParam` can only perform access by downcasting to that param type. +unsafe impl<'a, 'w, 's> SystemParamBuilder> for DynParamBuilder<'a> { + fn build( + self, + world: &mut World, + meta: &mut SystemMeta, + ) -> as SystemParam>::State { + (self.0)(world, meta) + } +} + /// A [`SystemParamBuilder`] for a [`Local`]. /// The provided value will be used as the initial value of the `Local`. pub struct LocalBuilder(pub T); @@ -312,6 +340,7 @@ unsafe impl<'s, T: FromWorld + Send + 'static> SystemParamBuilder> #[cfg(test)] mod tests { use crate as bevy_ecs; + use crate::entity::Entities; use crate::prelude::{Component, Query}; use crate::system::{Local, RunSystemOnce}; @@ -452,4 +481,35 @@ mod tests { let result = world.run_system_once(system); assert_eq!(result, 5); } + + #[test] + fn dyn_vec_builder() { + let mut world = World::new(); + + world.spawn(A); + world.spawn_empty(); + + let system = (vec![ + DynParamBuilder::new(LocalBuilder(3_usize)), + DynParamBuilder::new::>(QueryParamBuilder::new(|builder| { + builder.with::(); + })), + DynParamBuilder::new::<&Entities>(ParamBuilder), + ],) + .build_state(&mut world) + .build_system(|mut params: Vec| { + let local = *params[0].downcast_mut::>().unwrap(); + let query_count = params[1] + .downcast_mut::>() + .unwrap() + .iter() + .count(); + let _entities = params[2].downcast_mut::<&Entities>().unwrap(); + assert!(params[0].downcast_mut::>().is_none()); + local + query_count + }); + + let result = world.run_system_once(system); + assert_eq!(result, 4); + } } diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs index 105f50f300e91..50c4f979d8172 100644 --- a/crates/bevy_ecs/src/system/system_param.rs +++ b/crates/bevy_ecs/src/system/system_param.rs @@ -18,6 +18,7 @@ pub use bevy_ecs_macros::SystemParam; use bevy_ptr::UnsafeCellDeref; use bevy_utils::{all_tuples, synccell::SyncCell}; use std::{ + any::Any, fmt::Debug, marker::PhantomData, ops::{Deref, DerefMut}, @@ -1709,6 +1710,171 @@ unsafe impl SystemParam for PhantomData { // SAFETY: No world access. unsafe impl ReadOnlySystemParam for PhantomData {} +/// A [`SystemParam`] with a type that can be configured at runtime. +/// To be useful, this must be configured using a [`DynParamBuilder`](crate::system::DynParamBuilder) to build the system using a [`SystemParamBuilder`](crate::prelude::SystemParamBuilder). +// +// # SAFETY +// - The passed [`UnsafeWorldCell`] must have access to any world data +// registered in [`init_state`](SystemParam::init_state) for the +// inner system param, and the access must be valid for `'w`. +// - `world` must be the same `World` that was used to initialize [`state`](SystemParam::init_state). +pub struct DynSystemParam<'w, 's> { + /// A `ParamState` wrapping the state for the underlying system param. + state: &'s mut dyn Any, + world: UnsafeWorldCell<'w>, + system_meta: SystemMeta, + change_tick: Tick, +} + +impl<'w, 's> DynSystemParam<'w, 's> { + /// Returns `true` if the inner system param is the same as `T`. + pub fn is(&self) -> bool { + self.state.is::>() + } + + /// Returns the inner system param if it is the correct type. + /// This consumes the dyn param, so the returned param can have its original world and state lifetimes. + pub fn downcast(self) -> Option> { + // SAFETY: The `DynSystemParam` ensures the world matches and has access required by the inner system param. + // This consumes the `DynSystemParam`, so it is the only use of `world` with this access and it is available for `'w`. + unsafe { downcast::(self.state, &self.system_meta, self.world, self.change_tick) } + } + + /// Returns the inner system parameter if it is the correct type. + /// This borrows the dyn param, so the returned param is only valid for the duration of that borrow. + pub fn downcast_mut(&mut self) -> Option> { + // SAFETY: The `DynSystemParam` ensures the world matches and has access required by the inner system param. + // the world matched and had access required by the inner system param. + // This exclusively borrows the `DynSystemParam` for `'_`, so it is the only use of `world` with this access for `'_`. + unsafe { downcast::(self.state, &self.system_meta, self.world, self.change_tick) } + } + + /// Returns the inner system parameter if it is the correct type. + /// This borrows the dyn param, so the returned param is only valid for the duration of that borrow, + /// but since it only performs read access it can keep the original world lifetime. + /// This can be useful with methods like [`Query::iter_inner()`] or [`Res::into_inner()`] + /// to obtain references with the original world lifetime. + pub fn downcast_mut_inner( + &mut self, + ) -> Option> { + // SAFETY: The `DynSystemParam` ensures the world matches and has access required by the inner system param. + // the world matched and had access required by the inner system param. + // The inner system param only performs read access, so it's safe to copy that access for the full `'w` lifetime. + unsafe { downcast::(self.state, &self.system_meta, self.world, self.change_tick) } + } +} + +/// # SAFETY +/// - The passed [`UnsafeWorldCell`] must have access to any world data +/// registered in [`init_state`](SystemParam::init_state) for the +/// inner system param, and the access must be valid for `'w`. +/// - `world` must be the same `World` that was used to initialize [`state`](SystemParam::init_state). +unsafe fn downcast<'w, 's, T: SystemParam + 'static>( + state: &'s mut dyn Any, + system_meta: &SystemMeta, + world: UnsafeWorldCell<'w>, + change_tick: Tick, +) -> Option> { + state.downcast_mut::>().map(|state| { + // SAFETY: + // - The caller ensures the world has access for the underlying system param, + // and since the downcast succeeded, the underlying system param is T. + // - The caller ensures the `world` matches. + unsafe { T::get_param(&mut state.0, system_meta, world, change_tick) } + }) +} + +/// The [`SystemParam::State`] for a [`DynSystemParam`]. +pub struct DynSystemParamState(Box); + +impl DynSystemParamState { + pub(crate) fn new(state: T::State) -> Self { + Self(Box::new(ParamState::(state))) + } +} + +/// Allows a [`SystemParam::State`] to be used as a trait object for implementing [`DynSystemParam`]. +/// +/// # Safety +/// +/// This must only be implemented on `ParamState`, and `as_any_mut` must return `self`. +unsafe trait DynParamState: Sync + Send { + /// Casts the underlying `ParamState` to an `Any` so it can be downcast. + fn as_any_mut(&mut self) -> &mut dyn Any; + + /// For the specified [`Archetype`], registers the components accessed by this [`SystemParam`] (if applicable).a + /// + /// # Safety + /// `archetype` must be from the [`World`] used to initialize `state` in `init_state`. + unsafe fn new_archetype(&mut self, archetype: &Archetype, system_meta: &mut SystemMeta); + + /// Applies any deferred mutations stored in this [`SystemParam`]'s state. + /// This is used to apply [`Commands`] during [`apply_deferred`](crate::prelude::apply_deferred). + /// + /// [`Commands`]: crate::prelude::Commands + fn apply(&mut self, system_meta: &SystemMeta, world: &mut World); +} + +/// A wrapper around a [`SystemParam::State`] that can be used a a trait object in a [`DynSystemParam`]. +struct ParamState(T::State); + +// SAFETY: This is implemented on `ParamState`, and `as_any_mut` returns `self`. +unsafe impl DynParamState for ParamState { + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + unsafe fn new_archetype(&mut self, archetype: &Archetype, system_meta: &mut SystemMeta) { + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`. + unsafe { T::new_archetype(&mut self.0, archetype, system_meta) }; + } + + fn apply(&mut self, system_meta: &SystemMeta, world: &mut World) { + T::apply(&mut self.0, system_meta, world); + } +} + +// SAFETY: `init_state` creates a state of (), which performs no access. The interesting safety checks are on the `SystemParamBuilder`. +unsafe impl SystemParam for DynSystemParam<'_, '_> { + type State = DynSystemParamState; + + type Item<'world, 'state> = DynSystemParam<'world, 'state>; + + fn init_state(_world: &mut World, _system_meta: &mut SystemMeta) -> Self::State { + DynSystemParamState::new::<()>(()) + } + + unsafe fn get_param<'world, 'state>( + state: &'state mut Self::State, + system_meta: &SystemMeta, + world: UnsafeWorldCell<'world>, + change_tick: Tick, + ) -> Self::Item<'world, 'state> { + // SAFETY: + // - `state.0` is the underlying system param used to initialize our state, and `as_any_mut` just casts it to `Any`. + // - The caller ensures that the provided world is the same and has the required access. + DynSystemParam { + state: state.0.as_any_mut(), + world, + system_meta: system_meta.clone(), + change_tick, + } + } + + unsafe fn new_archetype( + state: &mut Self::State, + archetype: &Archetype, + system_meta: &mut SystemMeta, + ) { + // SAFETY: The caller ensures that `archetype` is from the World the state was initialized from in `init_state`. + unsafe { state.0.new_archetype(archetype, system_meta) }; + } + + fn apply(state: &mut Self::State, system_meta: &SystemMeta, world: &mut World) { + state.0.apply(system_meta, world); + } +} + #[cfg(test)] mod tests { use super::*;