diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs index 632501a56917f..2f4b868bd2045 100644 --- a/crates/bevy_ecs/src/system/system_param.rs +++ b/crates/bevy_ecs/src/system/system_param.rs @@ -596,6 +596,71 @@ impl<'a, T: Component + FromWorld> SystemParamFetch<'a> for LocalState { } } +/// A version of [`Local`] which requires manual initialization +pub struct Required<'a, T: Component>(&'a mut T); + +// SAFE: Required only accesses internal state +unsafe impl ReadOnlySystemParamFetch for RequiredState {} + +impl<'a, T: Component> Debug for Required<'a, T> +where + T: Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("Required").field(&self.0).finish() + } +} + +impl<'a, T: Component> Deref for Required<'a, T> { + type Target = T; + + #[inline] + fn deref(&self) -> &Self::Target { + self.0 + } +} + +impl<'a, T: Component> DerefMut for Required<'a, T> { + #[inline] + fn deref_mut(&mut self) -> &mut Self::Target { + self.0 + } +} + +/// The [`SystemParamState`] of [`Required`]. +pub struct RequiredState(T); + +impl<'a, T: Component> SystemParam for Required<'a, T> { + type Fetch = RequiredState; +} + +// SAFE: only local state is accessed +unsafe impl SystemParamState for RequiredState { + type Config = Option; + + fn init(_world: &mut World, _system_meta: &mut SystemMeta, config: Self::Config) -> Self { + Self(config.expect("Required must be initialized using config!")) + } + + fn default_config() -> Option { + None + } +} + +impl<'a, T: Component> SystemParamFetch<'a> for RequiredState { + type Item = Required<'a, T>; + + #[inline] + unsafe fn get_param( + state: &'a mut Self, + _system_meta: &SystemMeta, + _world: &'a World, + _change_tick: u32, + ) -> Self::Item { + Required(&mut state.0) + } +} + /// A [`SystemParam`] that grants access to the entities that had their `T` [`Component`] removed. /// /// # Examples