From 335f2903d99e44b7c5f3db061799033978d55523 Mon Sep 17 00:00:00 2001 From: Chris Russell <8494645+chescock@users.noreply.github.com> Date: Sun, 25 Aug 2024 10:23:44 -0400 Subject: [PATCH] SystemParamBuilder - Support dynamic system parameters (#14817) # Objective Support building systems with parameters whose types can be determined at runtime. ## Solution Create a `DynSystemParam` type that can be built using a `SystemParamBuilder` of any type and then downcast to the appropriate type dynamically. ## Example ```rust let system = ( DynParamBuilder::new(LocalBuilder(3_usize)), DynParamBuilder::new::>(QueryParamBuilder::new(|builder| { builder.with::(); })), DynParamBuilder::new::<&Entities>(ParamBuilder), ) .build_state(&mut world) .build_system( |mut p0: DynSystemParam, mut p1: DynSystemParam, mut p2: DynSystemParam| { let local = p0.downcast_mut::>().unwrap(); let query_count = p1.downcast_mut::>().unwrap(); let entities = p2.downcast_mut::<&Entities>().unwrap(); }, ); ``` --------- Co-authored-by: Alice Cecile Co-authored-by: Periwink --- crates/bevy_ecs/src/system/builder.rs | 60 ++++- crates/bevy_ecs/src/system/system_param.rs | 248 +++++++++++++++++++++ 2 files changed, 307 insertions(+), 1 deletion(-) diff --git a/crates/bevy_ecs/src/system/builder.rs b/crates/bevy_ecs/src/system/builder.rs index 2cecceeb3c2c0..861f915a70145 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}, @@ -251,6 +251,34 @@ macro_rules! impl_param_set_builder_tuple { all_tuples!(impl_param_set_builder_tuple, 1, 8, P, B, meta); +/// 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); @@ -271,6 +299,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}; @@ -382,4 +411,33 @@ mod tests { let result = world.run_system_once(system); assert_eq!(result, 5); } + + #[test] + fn dyn_builder() { + let mut world = World::new(); + + world.spawn(A); + world.spawn_empty(); + + let system = ( + DynParamBuilder::new(LocalBuilder(3_usize)), + DynParamBuilder::new::>(QueryParamBuilder::new(|builder| { + builder.with::(); + })), + DynParamBuilder::new::<&Entities>(ParamBuilder), + ) + .build_state(&mut world) + .build_system( + |mut p0: DynSystemParam, mut p1: DynSystemParam, mut p2: DynSystemParam| { + let local = *p0.downcast_mut::>().unwrap(); + let query_count = p1.downcast_mut::>().unwrap().iter().count(); + let _entities = p2.downcast_mut::<&Entities>().unwrap(); + assert!(p0.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 acc876a556b3c..7968338764b14 100644 --- a/crates/bevy_ecs/src/system/system_param.rs +++ b/crates/bevy_ecs/src/system/system_param.rs @@ -22,6 +22,7 @@ use bevy_utils::{all_tuples, synccell::SyncCell}; #[cfg(feature = "track_change_detection")] use std::panic::Location; use std::{ + any::Any, fmt::Debug, marker::PhantomData, ops::{Deref, DerefMut}, @@ -1674,6 +1675,253 @@ 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). +/// +/// # Examples +/// +/// ``` +/// # use bevy_ecs::{prelude::*, system::*}; +/// # +/// # #[derive(Default, Resource)] +/// # struct A; +/// # +/// # #[derive(Default, Resource)] +/// # struct B; +/// # +/// let mut world = World::new(); +/// world.init_resource::(); +/// world.init_resource::(); +/// +/// // If the inner parameter doesn't require any special building, use `ParamBuilder`. +/// // Either specify the type parameter on `DynParamBuilder::new()` ... +/// let system = (DynParamBuilder::new::>(ParamBuilder),) +/// .build_state(&mut world) +/// .build_system(expects_res_a); +/// world.run_system_once(system); +/// +/// // ... or use a factory method on `ParamBuilder` that returns a specific type. +/// let system = (DynParamBuilder::new(ParamBuilder::resource::()),) +/// .build_state(&mut world) +/// .build_system(expects_res_a); +/// world.run_system_once(system); +/// +/// fn expects_res_a(mut param: DynSystemParam) { +/// // Use the `downcast` methods to retrieve the inner parameter. +/// // They will return `None` if the type does not match. +/// assert!(param.is::>()); +/// assert!(!param.is::>()); +/// assert!(param.downcast_mut::>().is_none()); +/// let foo: Res = param.downcast::>().unwrap(); +/// } +/// +/// let system = ( +/// // If the inner parameter also requires building, +/// // pass the appropriate `SystemParamBuilder`. +/// DynParamBuilder::new(LocalBuilder(10usize)), +/// // `DynSystemParam` is just an ordinary `SystemParam`, +/// // and can be combined with other parameters as usual! +/// ParamBuilder::query(), +/// ) +/// .build_state(&mut world) +/// .build_system(|param: DynSystemParam, query: Query<()>| { +/// let local: Local = param.downcast::>().unwrap(); +/// assert_eq!(*local, 10); +/// }); +/// world.run_system_once(system); +/// ``` +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> { + /// # SAFETY + /// - `state` must be a `ParamState` for some inner `T: SystemParam`. + /// - The passed [`UnsafeWorldCell`] must have access to any world data registered + /// in [`init_state`](SystemParam::init_state) for the inner system param. + /// - `world` must be the same `World` that was used to initialize + /// [`state`](SystemParam::init_state) for the inner system param. + unsafe fn new( + state: &'s mut dyn Any, + world: UnsafeWorldCell<'w>, + system_meta: SystemMeta, + change_tick: Tick, + ) -> Self { + Self { + state, + world, + system_meta, + change_tick, + } + } + + /// 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: + // - `DynSystemParam::new()` ensures `state` is a `ParamState`, that the world matches, + // and that it 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: + // - `DynSystemParam::new()` ensures `state` is a `ParamState`, that the world matches, + // and that it has 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: + // - `DynSystemParam::new()` ensures `state` is a `ParamState`, that the world matches, + // and that it has 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 +/// - `state` must be a `ParamState` for some inner `T: SystemParam`. +/// - The passed [`UnsafeWorldCell`] must have access to any world data registered +/// in [`init_state`](SystemParam::init_state) for the inner system param. +/// - `world` must be the same `World` that was used to initialize +/// [`state`](SystemParam::init_state) for the inner system param. +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`]. +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); + + /// Queues any deferred mutations to be applied at the next [`apply_deferred`](crate::prelude::apply_deferred). + fn queue(&mut self, system_meta: &SystemMeta, world: DeferredWorld); +} + +/// A wrapper around a [`SystemParam::State`] that can be used as a trait object in a [`DynSystemParam`]. +struct ParamState(T::State); + +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); + } + + fn queue(&mut self, system_meta: &SystemMeta, world: DeferredWorld) { + T::queue(&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 a boxed `ParamState`, and its implementation of `as_any_mut` returns `self`. + // - The state was obtained from `SystemParamBuilder::build()`, which registers all [`World`] accesses used + // by [`SystemParam::get_param`] with the provided [`system_meta`](SystemMeta). + // - The caller ensures that the provided world is the same and has the required access. + unsafe { + DynSystemParam::new( + state.0.as_any_mut(), + world, + 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); + } + + fn queue(state: &mut Self::State, system_meta: &SystemMeta, world: DeferredWorld) { + state.0.queue(system_meta, world); + } +} + #[cfg(test)] mod tests { use super::*;