Refactor FunctionSystem to use a single Option

crates/bevy_ecs/src/system/function_system.rs

@@ -395,9 +395,11 @@ impl<Param: SystemParam> SystemState<Param> {
     ) -> FunctionSystem<Marker, F> {
         FunctionSystem {
             func,
-            param_state: Some(self.param_state),
+            state: Some(FunctionSystemState {
+                param: self.param_state,
+                world_id: self.world_id,
+            }),
             system_meta: self.meta,
-            world_id: Some(self.world_id),
             archetype_generation: self.archetype_generation,
             marker: PhantomData,
         }
@@ -602,14 +604,25 @@ where
     F: SystemParamFunction<Marker>,
 {
     func: F,
-    pub(crate) param_state: Option<<F::Param as SystemParam>::State>,
-    pub(crate) system_meta: SystemMeta,
-    world_id: Option<WorldId>,
+    state: Option<FunctionSystemState<F::Param>>,
+    system_meta: SystemMeta,
     archetype_generation: ArchetypeGeneration,
     // NOTE: PhantomData<fn()-> T> gives this safe Send/Sync impls
     marker: PhantomData<fn() -> Marker>,
 }
 
+/// The state of a [`FunctionSystem`], which must be initialized with
+/// [`System::initialize`] before the system can be run. A panic will occur if
+/// the system is run without being initialized.
+struct FunctionSystemState<P: SystemParam> {
+    /// The cached state of the system's [`SystemParam`]s.
+    param: P::State,
+    /// The id of the [`World`] this system was initialized with. If the world
+    /// passed to [`System::update_archetype_component_access`] does not match
+    /// this id, a panic will occur.
+    world_id: WorldId,
+}
+
 impl<Marker, F> FunctionSystem<Marker, F>
 where
     F: SystemParamFunction<Marker>,
@@ -631,9 +644,8 @@ where
     fn clone(&self) -> Self {
         Self {
             func: self.func.clone(),
-            param_state: None,
+            state: None,
             system_meta: SystemMeta::new::<F>(),
-            world_id: None,
             archetype_generation: ArchetypeGeneration::initial(),
             marker: PhantomData,
         }
@@ -653,9 +665,8 @@ where
     fn into_system(func: Self) -> Self::System {
         FunctionSystem {
             func,
-            param_state: None,
+            state: None,
             system_meta: SystemMeta::new::<F>(),
-            world_id: None,
             archetype_generation: ArchetypeGeneration::initial(),
             marker: PhantomData,
         }
@@ -669,7 +680,8 @@ where
     /// Message shown when a system isn't initialized
     // When lines get too long, rustfmt can sometimes refuse to format them.
     // Work around this by storing the message separately.
-    const PARAM_MESSAGE: &'static str = "System's param_state was not found. Did you forget to initialize this system before running it?";
+    const ERROR_UNINITIALIZED: &'static str =
+        "System's state was not found. Did you forget to initialize this system before running it?";
 }
 
 impl<Marker, F> System for FunctionSystem<Marker, F>
@@ -721,39 +733,34 @@ where
 
         let change_tick = world.increment_change_tick();
 
+        let param_state = &mut self.state.as_mut().expect(Self::ERROR_UNINITIALIZED).param;
         // SAFETY:
         // - The caller has invoked `update_archetype_component_access`, which will panic
         //   if the world does not match.
         // - All world accesses used by `F::Param` have been registered, so the caller
         //   will ensure that there are no data access conflicts.
-        let params = unsafe {
-            F::Param::get_param(
-                self.param_state.as_mut().expect(Self::PARAM_MESSAGE),
-                &self.system_meta,
-                world,
-                change_tick,
-            )
-        };
+        let params =
+            unsafe { F::Param::get_param(param_state, &self.system_meta, world, change_tick) };
         let out = self.func.run(input, params);
         self.system_meta.last_run = change_tick;
         out
     }
 
     #[inline]
     fn apply_deferred(&mut self, world: &mut World) {
-        let param_state = self.param_state.as_mut().expect(Self::PARAM_MESSAGE);
+        let param_state = &mut self.state.as_mut().expect(Self::ERROR_UNINITIALIZED).param;
         F::Param::apply(param_state, &self.system_meta, world);
     }
 
     #[inline]
     fn queue_deferred(&mut self, world: DeferredWorld) {
-        let param_state = self.param_state.as_mut().expect(Self::PARAM_MESSAGE);
+        let param_state = &mut self.state.as_mut().expect(Self::ERROR_UNINITIALIZED).param;
         F::Param::queue(param_state, &self.system_meta, world);
     }
 
     #[inline]
     unsafe fn validate_param_unsafe(&mut self, world: UnsafeWorldCell) -> bool {
-        let param_state = self.param_state.as_ref().expect(Self::PARAM_MESSAGE);
+        let param_state = &self.state.as_ref().expect(Self::ERROR_UNINITIALIZED).param;
         // SAFETY:
         // - The caller has invoked `update_archetype_component_access`, which will panic
         //   if the world does not match.
@@ -768,29 +775,32 @@ where
 
     #[inline]
     fn initialize(&mut self, world: &mut World) {
-        if let Some(id) = self.world_id {
+        if let Some(state) = &self.state {
             assert_eq!(
-                id,
+                state.world_id,
                 world.id(),
                 "System built with a different world than the one it was added to.",
             );
         } else {
-            self.world_id = Some(world.id());
-            self.param_state = Some(F::Param::init_state(world, &mut self.system_meta));
+            self.state = Some(FunctionSystemState {
+                param: F::Param::init_state(world, &mut self.system_meta),
+                world_id: world.id(),
+            });
         }
         self.system_meta.last_run = world.change_tick().relative_to(Tick::MAX);
     }
 
     fn update_archetype_component_access(&mut self, world: UnsafeWorldCell) {
-        assert_eq!(self.world_id, Some(world.id()), "Encountered a mismatched World. A System cannot be used with Worlds other than the one it was initialized with.");
+        let state = self.state.as_mut().expect(Self::ERROR_UNINITIALIZED);
+        assert_eq!(state.world_id, world.id(), "Encountered a mismatched World. A System cannot be used with Worlds other than the one it was initialized with.");
+
         let archetypes = world.archetypes();
         let old_generation =
             core::mem::replace(&mut self.archetype_generation, archetypes.generation());
 
         for archetype in &archetypes[old_generation..] {
-            let param_state = self.param_state.as_mut().unwrap();
             // SAFETY: The assertion above ensures that the param_state was initialized from `world`.
-            unsafe { F::Param::new_archetype(param_state, archetype, &mut self.system_meta) };
+            unsafe { F::Param::new_archetype(&mut state.param, archetype, &mut self.system_meta) };
         }
     }