Support dynamic system parameters.

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<P>>
     }
 }
 
+/// 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);
@@ -312,6 +340,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};
 
@@ -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::<Query<()>>(QueryParamBuilder::new(|builder| {
+                builder.with::<A>();
+            })),
+            DynParamBuilder::new::<&Entities>(ParamBuilder),
+        ],)
+            .build_state(&mut world)
+            .build_system(|mut params: Vec<DynSystemParam>| {
+                let local = *params[0].downcast_mut::<Local<usize>>().unwrap();
+                let query_count = params[1]
+                    .downcast_mut::<Query<()>>()
+                    .unwrap()
+                    .iter()
+                    .count();
+                let _entities = params[2].downcast_mut::<&Entities>().unwrap();
+                assert!(params[0].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

@@ -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<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).
+//
+// # 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<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> {
+    /// 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: 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::<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: 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::<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: 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::<T>(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<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`].
+///
+/// # Safety
+///
+/// This must only be implemented on `ParamState<T>`, and `as_any_mut` must return `self`.
+unsafe 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);
+}
+
+/// A wrapper around a [`SystemParam::State`] that can be used a a trait object in a [`DynSystemParam`].
+struct ParamState<T: SystemParam>(T::State);
+
+// SAFETY: This is implemented on `ParamState<T>`, and `as_any_mut` returns `self`.
+unsafe 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);
+    }
+}
+
+// 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::*;