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::<Query<()>>(QueryParamBuilder::new(|builder| {
        builder.with::<A>();
    })),
    DynParamBuilder::new::<&Entities>(ParamBuilder),
)
    .build_state(&mut world)
    .build_system(
        |mut p0: DynSystemParam, mut p1: DynSystemParam, mut p2: DynSystemParam| {
            let local = p0.downcast_mut::<Local<usize>>().unwrap();
            let query_count = p1.downcast_mut::<Query<()>>().unwrap();
            let entities = p2.downcast_mut::<&Entities>().unwrap();
        },
    );
```

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Periwink <charlesbour@gmail.com>

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<dyn FnOnce(&mut World, &mut SystemMeta) -> 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<T: SystemParam + 'static>(builder: impl SystemParamBuilder<T> + 'a) -> Self {
+ Self(Box::new(|world, meta| {
+ DynSystemParamState::new::<T>(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<DynSystemParam<'w, 's>> for DynParamBuilder<'a> {
+ fn build(
+ self,
+ world: &mut World,
+ meta: &mut SystemMeta,
+ ) -> <DynSystemParam<'w, 's> 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<T>(pub T);
@@ -271,6 +299,7 @@ unsafe impl<'s, T: FromWorld + Send + 'static> SystemParamBuilder<Local<'s, T>>
 #[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::<Query<()>>(QueryParamBuilder::new(|builder| {
+ builder.with::<A>();
+ })),
+ DynParamBuilder::new::<&Entities>(ParamBuilder),
+ )
+ .build_state(&mut world)
+ .build_system(
+ |mut p0: DynSystemParam, mut p1: DynSystemParam, mut p2: DynSystemParam| {
+ let local = *p0.downcast_mut::<Local<usize>>().unwrap();
+ let query_count = p1.downcast_mut::<Query<()>>().unwrap().iter().count();
+ let _entities = p2.downcast_mut::<&Entities>().unwrap();
+ assert!(p0.downcast_mut::<Query<()>>().is_none());
+ local + query_count
+ },
+ );
+
+ let result = world.run_system_once(system);
+ assert_eq!(result, 4);
+ }
 }

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<T: ?Sized> SystemParam for PhantomData<T> {
 // SAFETY: No world access.
 unsafe impl<T: ?Sized> ReadOnlySystemParam for PhantomData<T> {}
 
+/// 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::<A>();
+/// world.init_resource::<B>();
+///
+/// // If the inner parameter doesn't require any special building, use `ParamBuilder`.
+/// // Either specify the type parameter on `DynParamBuilder::new()` ...
+/// let system = (DynParamBuilder::new::<Res<A>>(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::<A>()),)
+/// .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::<Res<A>>());
+/// assert!(!param.is::<Res<B>>());
+/// assert!(param.downcast_mut::<Res<B>>().is_none());
+/// let foo: Res<A> = param.downcast::<Res<A>>().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<usize> = param.downcast::<Local<usize>>().unwrap();
+/// assert_eq!(*local, 10);
+/// });
+/// world.run_system_once(system);
+/// ```
+pub struct DynSystemParam<'w, 's> {
+ /// A `ParamState<T>` 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<T>` 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<T: SystemParam + 'static>(&self) -> bool {
+ self.state.is::<ParamState<T>>()
+ }
+
+ /// 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<T: SystemParam + 'static>(self) -> Option<T::Item<'w, 's>> {
+ // SAFETY:
+ // - `DynSystemParam::new()` ensures `state` is a `ParamState<T>`, 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::<T>(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<T: SystemParam + 'static>(&mut self) -> Option<T::Item<'_, '_>> {
+ // SAFETY:
+ // - `DynSystemParam::new()` ensures `state` is a `ParamState<T>`, 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::<T>(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<T: ReadOnlySystemParam + 'static>(
+ &mut self,
+ ) -> Option<T::Item<'w, '_>> {
+ // SAFETY:
+ // - `DynSystemParam::new()` ensures `state` is a `ParamState<T>`, 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::<T>(self.state, &self.system_meta, self.world, self.change_tick) }
+ }
+}
+
+/// # SAFETY
+/// - `state` must be a `ParamState<T>` 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<T::Item<'w, 's>> {
+ state.downcast_mut::<ParamState<T>>().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<dyn DynParamState>);
+
+impl DynSystemParamState {
+ pub(crate) fn new<T: SystemParam + 'static>(state: T::State) -> Self {
+ Self(Box::new(ParamState::<T>(state)))
+ }
+}
+
+/// Allows a [`SystemParam::State`] to be used as a trait object for implementing [`DynSystemParam`].
+trait DynParamState: Sync + Send {
+ /// Casts the underlying `ParamState<T>` 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: SystemParam>(T::State);
+
+impl<T: SystemParam + 'static> DynParamState for ParamState<T> {
+ 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<T>`, 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::*;