diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs index e98a504ca27bd..9dcdb9f4d2201 100644 --- a/crates/bevy_ecs/src/system/system_param.rs +++ b/crates/bevy_ecs/src/system/system_param.rs @@ -147,6 +147,11 @@ pub struct QuerySetState(T); impl_query_set!(); +/// Shared borrow of a Resource +/// +/// When used as a system parameter, panics if resource does not exist. +/// +/// Use `Option>` if the resource might not always exist. pub struct Res<'w, T> { value: &'w T, flags: ComponentFlags, @@ -235,6 +240,42 @@ impl<'a, T: Component> SystemParamFetch<'a> for ResState { } } +pub struct OptionResState(ResState); + +impl<'a, T: Component> SystemParam for Option> { + type Fetch = OptionResState; +} + +unsafe impl SystemParamState for OptionResState { + type Config = (); + fn init(world: &mut World, system_state: &mut SystemState, _config: Self::Config) -> Self { + Self(ResState::init(world, system_state, ())) + } +} + +impl<'a, T: Component> SystemParamFetch<'a> for OptionResState { + type Item = Option>; + + #[inline] + unsafe fn get_param( + state: &'a mut Self, + _system_state: &'a SystemState, + world: &'a World, + ) -> Self::Item { + world + .get_populated_resource_column(state.0.component_id) + .map(|column| Res { + value: &*column.get_ptr().as_ptr().cast::(), + flags: *column.get_flags_mut_ptr(), + }) + } +} + +/// Unique borrow of a Resource +/// +/// When used as a system parameter, panics if resource does not exist. +/// +/// Use `Option>` if the resource might not always exist. pub struct ResMut<'w, T> { value: &'w mut T, flags: &'w mut ComponentFlags, @@ -334,6 +375,37 @@ impl<'a, T: Component> SystemParamFetch<'a> for ResMutState { } } +pub struct OptionResMutState(ResMutState); + +impl<'a, T: Component> SystemParam for Option> { + type Fetch = OptionResMutState; +} + +unsafe impl SystemParamState for OptionResMutState { + type Config = (); + fn init(world: &mut World, system_state: &mut SystemState, _config: Self::Config) -> Self { + Self(ResMutState::init(world, system_state, ())) + } +} + +impl<'a, T: Component> SystemParamFetch<'a> for OptionResMutState { + type Item = Option>; + + #[inline] + unsafe fn get_param( + state: &'a mut Self, + _system_state: &'a SystemState, + world: &'a World, + ) -> Self::Item { + world + .get_resource_unchecked_mut_with_id(state.0.component_id) + .map(|value| ResMut { + value: value.value, + flags: value.flags, + }) + } +} + impl<'a> SystemParam for Commands<'a> { type Fetch = CommandQueue; }