From 13aef0503859925425e11cb9785ac0bb623b83e9 Mon Sep 17 00:00:00 2001
From: Jasen Borisov <jamadazi@protonmail.ch>
Date: Mon, 8 Mar 2021 20:12:22 +0000
Subject: [PATCH] impl SystemParam for Option<Res<T>> / Option<ResMut<T>>
 (#1494)

This allows users to write systems that do not panic if a resource does not exist at runtime (such as if it has not been inserted yet).

This is a copy-paste of the impls for `Res` and `ResMut`, with an extra check to see if the resource exists.

There might be a cleaner way to do it than this check. I don't know.
---
 crates/bevy_ecs/src/system/system_param.rs | 72 ++++++++++++++++++++++
 1 file changed, 72 insertions(+)

diff --git a/crates/bevy_ecs/src/system/system_param.rs b/crates/bevy_ecs/src/system/system_param.rs
index e98a504ca2..9dcdb9f4d2 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>(T);
 
 impl_query_set!();
 
+/// Shared borrow of a Resource
+///
+/// When used as a system parameter, panics if resource does not exist.
+///
+/// Use `Option<Res<T>>` 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<T> {
     }
 }
 
+pub struct OptionResState<T>(ResState<T>);
+
+impl<'a, T: Component> SystemParam for Option<Res<'a, T>> {
+    type Fetch = OptionResState<T>;
+}
+
+unsafe impl<T: Component> SystemParamState for OptionResState<T> {
+    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<T> {
+    type Item = Option<Res<'a, T>>;
+
+    #[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::<T>(),
+                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<ResMut<T>>` 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<T> {
     }
 }
 
+pub struct OptionResMutState<T>(ResMutState<T>);
+
+impl<'a, T: Component> SystemParam for Option<ResMut<'a, T>> {
+    type Fetch = OptionResMutState<T>;
+}
+
+unsafe impl<T: Component> SystemParamState for OptionResMutState<T> {
+    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<T> {
+    type Item = Option<ResMut<'a, T>>;
+
+    #[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;
 }