Add World::try_run_schedule (#8028)

Co-authored-by: James Liu <contact@jamessliu.com>

crates/bevy_app/src/app.rs

@@ -358,15 +358,10 @@ impl App {
                     .run_if(in_state(variant)),
             );
         }
-        // These are different for loops to avoid conflicting access to self
-        for variant in S::variants() {
-            if self.get_schedule(OnEnter(variant.clone())).is_none() {
-                self.add_schedule(OnEnter(variant.clone()), Schedule::new());
-            }
-            if self.get_schedule(OnExit(variant.clone())).is_none() {
-                self.add_schedule(OnExit(variant), Schedule::new());
-            }
-        }
+
+        // The OnEnter, OnExit, and OnTransition schedules are lazily initialized
+        // (i.e. when the first system is added to them), and World::try_run_schedule is used to fail
+        // gracefully if they aren't present.
 
         self
     }

crates/bevy_ecs/Cargo.toml

@@ -27,6 +27,7 @@ fixedbitset = "0.4.2"
 rustc-hash = "1.1"
 downcast-rs = "1.2"
 serde = { version = "1", features = ["derive"] }
+thiserror = "1.0"
 
 [dev-dependencies]
 rand = "0.8"

crates/bevy_ecs/src/schedule/state.rs

@@ -4,7 +4,7 @@ use std::mem;
 
 use crate as bevy_ecs;
 use crate::change_detection::DetectChangesMut;
-use crate::schedule::{ScheduleLabel, Schedules, SystemSet};
+use crate::schedule::{ScheduleLabel, SystemSet};
 use crate::system::Resource;
 use crate::world::World;
 
@@ -100,15 +100,18 @@ impl<S: States> NextState<S> {
     }
 }
 
-/// Run the enter schedule for the current state
+/// Run the enter schedule (if it exists) for the current state.
 pub fn run_enter_schedule<S: States>(world: &mut World) {
-    world.run_schedule(OnEnter(world.resource::<State<S>>().0.clone()));
+    world
+        .try_run_schedule(OnEnter(world.resource::<State<S>>().0.clone()))
+        .ok();
 }
 
 /// If a new state is queued in [`NextState<S>`], this system:
 /// - Takes the new state value from [`NextState<S>`] and updates [`State<S>`].
-/// - Runs the [`OnExit(exited_state)`] schedule.
-/// - Runs the [`OnEnter(entered_state)`] schedule.
+/// - Runs the [`OnExit(exited_state)`] schedule, if it exists.
+/// - Runs the [`OnTransition { from: exited_state, to: entered_state }`](OnTransition), if it exists.
+/// - Runs the [`OnEnter(entered_state)`] schedule, if it exists.
 pub fn apply_state_transition<S: States>(world: &mut World) {
     // We want to take the `NextState` resource,
     // but only mark it as changed if it wasn't empty.
@@ -117,16 +120,15 @@ pub fn apply_state_transition<S: States>(world: &mut World) {
         next_state_resource.set_changed();
 
         let exited = mem::replace(&mut world.resource_mut::<State<S>>().0, entered.clone());
-        world.run_schedule(OnExit(exited.clone()));
 
-        let transition_schedule = OnTransition {
-            from: exited,
-            to: entered.clone(),
-        };
-        if world.resource::<Schedules>().contains(&transition_schedule) {
-            world.run_schedule(transition_schedule);
-        }
-
-        world.run_schedule(OnEnter(entered));
+        // Try to run the schedules if they exist.
+        world.try_run_schedule(OnExit(exited.clone())).ok();
+        world
+            .try_run_schedule(OnTransition {
+                from: exited,
+                to: entered.clone(),
+            })
+            .ok();
+        world.try_run_schedule(OnEnter(entered)).ok();
     }
 }

crates/bevy_ecs/src/world/error.rs

@@ -0,0 +1,10 @@
+use thiserror::Error;
+
+use crate::schedule::BoxedScheduleLabel;
+
+/// The error type returned by [`World::try_run_schedule`] if the provided schedule does not exist.
+///
+/// [`World::try_run_schedule`]: crate::world::World::try_run_schedule
+#[derive(Error, Debug)]
+#[error("The schedule with the label {0:?} was not found.")]
+pub struct TryRunScheduleError(pub BoxedScheduleLabel);

crates/bevy_ecs/src/world/mod.rs

@@ -1,4 +1,5 @@
 mod entity_ref;
+pub mod error;
 mod spawn_batch;
 pub mod unsafe_world_cell;
 mod world_cell;
@@ -20,6 +21,7 @@ use crate::{
     schedule::{Schedule, ScheduleLabel, Schedules},
     storage::{ResourceData, Storages},
     system::Resource,
+    world::error::TryRunScheduleError,
 };
 use bevy_ptr::{OwningPtr, Ptr};
 use bevy_utils::tracing::warn;
@@ -1714,6 +1716,47 @@ impl World {
         schedules.insert(label, schedule);
     }
 
+    /// Attempts to run the [`Schedule`] associated with the `label` a single time,
+    /// and returns a [`TryRunScheduleError`] if the schedule does not exist.
+    ///
+    /// The [`Schedule`] is fetched from the [`Schedules`] resource of the world by its label,
+    /// and system state is cached.
+    ///
+    /// For simple testing use cases, call [`Schedule::run(&mut world)`](Schedule::run) instead.
+    pub fn try_run_schedule(
+        &mut self,
+        label: impl ScheduleLabel,
+    ) -> Result<(), TryRunScheduleError> {
+        self.try_run_schedule_ref(&label)
+    }
+
+    /// Attempts to run the [`Schedule`] associated with the `label` a single time,
+    /// and returns a [`TryRunScheduleError`] if the schedule does not exist.
+    ///
+    /// Unlike the `try_run_schedule` method, this method takes the label by reference, which can save a clone.
+    ///
+    /// The [`Schedule`] is fetched from the [`Schedules`] resource of the world by its label,
+    /// and system state is cached.
+    ///
+    /// For simple testing use cases, call [`Schedule::run(&mut world)`](Schedule::run) instead.
+    pub fn try_run_schedule_ref(
+        &mut self,
+        label: &dyn ScheduleLabel,
+    ) -> Result<(), TryRunScheduleError> {
+        let Some((extracted_label, mut schedule)) = self.resource_mut::<Schedules>().remove_entry(label) else {
+            return Err(TryRunScheduleError(label.dyn_clone()));
+        };
+
+        // TODO: move this span to Schedule::run
+        #[cfg(feature = "trace")]
+        let _span = bevy_utils::tracing::info_span!("schedule", name = ?extracted_label).entered();
+        schedule.run(self);
+        self.resource_mut::<Schedules>()
+            .insert(extracted_label, schedule);
+
+        Ok(())
+    }
+
     /// Runs the [`Schedule`] associated with the `label` a single time.
     ///
     /// The [`Schedule`] is fetched from the [`Schedules`] resource of the world by its label,
@@ -1741,17 +1784,8 @@ impl World {
     ///
     /// Panics if the requested schedule does not exist, or the [`Schedules`] resource was not added.
     pub fn run_schedule_ref(&mut self, label: &dyn ScheduleLabel) {
-        let (extracted_label, mut schedule) = self
-            .resource_mut::<Schedules>()
-            .remove_entry(label)
-            .unwrap_or_else(|| panic!("The schedule with the label {label:?} was not found."));
-
-        // TODO: move this span to Schedule::run
-        #[cfg(feature = "trace")]
-        let _span = bevy_utils::tracing::info_span!("schedule", name = ?extracted_label).entered();
-        schedule.run(self);
-        self.resource_mut::<Schedules>()
-            .insert(extracted_label, schedule);
+        self.try_run_schedule_ref(label)
+            .unwrap_or_else(|e| panic!("{}", e));
     }
 }