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 <[email protected]>
Co-authored-by: Periwink <[email protected]>

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::*;