Add a scope API for world schedules (#8387)

# Objective

If you want to execute a schedule on the world using arbitrarily complex
behavior, you currently need to use "hokey-pokey strats": remove the
schedule from the world, do your thing, and add it back to the world.
Not only is this cumbersome, it's potentially error-prone as one might
forget to re-insert the schedule.

## Solution

Add the `World::{try}schedule_scope{ref}` family of functions, which is
a convenient abstraction over hokey pokey strats. This method
essentially works the same way as `World::resource_scope`.

### Example

```rust
// Run the schedule five times.
world.schedule_scope(MySchedule, |world, schedule| {
    for _ in 0..5 {
        schedule.run(world);
    }
});
```

---

## Changelog

Added the `World::schedule_scope` family of methods, which provide a way
to get mutable access to a world and one of its schedules at the same
time.

---------

Co-authored-by: James Liu <[email protected]>

Author: joseph-gio

crates/bevy_ecs/src/world/mod.rs

@@ -1721,33 +1721,43 @@ 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.
+    /// Temporarily removes the schedule associated with `label` from the world,
+    /// runs user code, and finally re-adds the schedule.
+    /// This returns a [`TryRunScheduleError`] if there is no schedule
+    /// associated with `label`.
     ///
     /// 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(
+    /// For simple cases where you just need to call the schedule once,
+    /// consider using [`World::try_run_schedule`] instead.
+    /// For other use cases, see the example on [`World::schedule_scope`].
+    pub fn try_schedule_scope<R>(
         &mut self,
         label: impl ScheduleLabel,
-    ) -> Result<(), TryRunScheduleError> {
-        self.try_run_schedule_ref(&label)
+        f: impl FnOnce(&mut World, &mut Schedule) -> R,
+    ) -> Result<R, TryRunScheduleError> {
+        self.try_schedule_scope_ref(&label, f)
     }
 
-    /// Attempts to run the [`Schedule`] associated with the `label` a single time,
-    /// and returns a [`TryRunScheduleError`] if the schedule does not exist.
+    /// Temporarily removes the schedule associated with `label` from the world,
+    /// runs user code, and finally re-adds the schedule.
+    /// This returns a [`TryRunScheduleError`] if there is no schedule
+    /// associated with `label`.
     ///
     /// 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(
+    /// For simple cases where you just need to call the schedule once,
+    /// consider using [`World::try_run_schedule_ref`] instead.
+    /// For other use cases, see the example on [`World::schedule_scope`].
+    pub fn try_schedule_scope_ref<R>(
         &mut self,
         label: &dyn ScheduleLabel,
-    ) -> Result<(), TryRunScheduleError> {
+        f: impl FnOnce(&mut World, &mut Schedule) -> R,
+    ) -> Result<R, TryRunScheduleError> {
         let Some((extracted_label, mut schedule))
             = self.get_resource_mut::<Schedules>().and_then(|mut s| s.remove_entry(label))
         else {
@@ -1757,11 +1767,115 @@ impl World {
         // 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>()
+        let value = f(self, &mut schedule);
+
+        let old = self
+            .resource_mut::<Schedules>()
             .insert(extracted_label, schedule);
+        if old.is_some() {
+            warn!("Schedule `{label:?}` was inserted during a call to `World::schedule_scope`: its value has been overwritten");
+        }
+
+        Ok(value)
+    }
+
+    /// Temporarily removes the schedule associated with `label` from the world,
+    /// runs user code, and finally re-adds the schedule.
+    ///
+    /// The [`Schedule`] is fetched from the [`Schedules`] resource of the world by its label,
+    /// and system state is cached.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use bevy_ecs::{prelude::*, schedule::ScheduleLabel};
+    /// # #[derive(ScheduleLabel, Debug, Clone, Copy, PartialEq, Eq, Hash)]
+    /// # pub struct MySchedule;
+    /// # #[derive(Resource)]
+    /// # struct Counter(usize);
+    /// #
+    /// # let mut world = World::new();
+    /// # world.insert_resource(Counter(0));
+    /// # let mut schedule = Schedule::new();
+    /// # schedule.add_systems(tick_counter);
+    /// # world.init_resource::<Schedules>();
+    /// # world.add_schedule(schedule, MySchedule);
+    /// # fn tick_counter(mut counter: ResMut<Counter>) { counter.0 += 1; }
+    /// // Run the schedule five times.
+    /// world.schedule_scope(MySchedule, |world, schedule| {
+    ///     for _ in 0..5 {
+    ///         schedule.run(world);
+    ///     }
+    /// });
+    /// # assert_eq!(world.resource::<Counter>().0, 5);
+    /// ```
+    ///
+    /// For simple cases where you just need to call the schedule once,
+    /// consider using [`World::run_schedule`] instead.
+    ///
+    /// # Panics
+    ///
+    /// If the requested schedule does not exist.
+    pub fn schedule_scope<R>(
+        &mut self,
+        label: impl ScheduleLabel,
+        f: impl FnOnce(&mut World, &mut Schedule) -> R,
+    ) -> R {
+        self.schedule_scope_ref(&label, f)
+    }
 
-        Ok(())
+    /// Temporarily removes the schedule associated with `label` from the world,
+    /// runs user code, and finally re-adds the schedule.
+    ///
+    /// Unlike the `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 cases where you just need to call the schedule,
+    /// consider using [`World::run_schedule_ref`] instead.
+    /// For other use cases, see the example on [`World::schedule_scope`].
+    ///
+    /// # Panics
+    ///
+    /// If the requested schedule does not exist.
+    pub fn schedule_scope_ref<R>(
+        &mut self,
+        label: &dyn ScheduleLabel,
+        f: impl FnOnce(&mut World, &mut Schedule) -> R,
+    ) -> R {
+        self.try_schedule_scope_ref(label, f)
+            .unwrap_or_else(|e| panic!("{e}"))
+    }
+
+    /// 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> {
+        self.try_schedule_scope_ref(label, |world, sched| sched.run(world))
     }
 
     /// Runs the [`Schedule`] associated with the `label` a single time.
@@ -1773,7 +1887,7 @@ impl World {
     ///
     /// # Panics
     ///
-    /// Panics if the requested schedule does not exist, or the [`Schedules`] resource was not added.
+    /// If the requested schedule does not exist.
     pub fn run_schedule(&mut self, label: impl ScheduleLabel) {
         self.run_schedule_ref(&label);
     }
@@ -1789,10 +1903,9 @@ impl World {
     ///
     /// # Panics
     ///
-    /// Panics if the requested schedule does not exist, or the [`Schedules`] resource was not added.
+    /// If the requested schedule does not exist.
     pub fn run_schedule_ref(&mut self, label: &dyn ScheduleLabel) {
-        self.try_run_schedule_ref(label)
-            .unwrap_or_else(|e| panic!("{}", e));
+        self.schedule_scope_ref(label, |world, sched| sched.run(world));
     }
 }
 

crates/bevy_time/src/fixed_timestep.rs

@@ -107,16 +107,11 @@ pub fn run_fixed_update_schedule(world: &mut World) {
     fixed_time.tick(delta_time);
 
     // Run the schedule until we run out of accumulated time
-    let mut check_again = true;
-    while check_again {
-        let mut fixed_time = world.resource_mut::<FixedTime>();
-        let fixed_time_run = fixed_time.expend().is_ok();
-        if fixed_time_run {
-            let _ = world.try_run_schedule(FixedUpdate);
-        } else {
-            check_again = false;
+    let _ = world.try_schedule_scope(FixedUpdate, |world, schedule| {
+        while world.resource_mut::<FixedTime>().expend().is_ok() {
+            schedule.run(world);
         }
-    }
+    });
 }
 
 #[cfg(test)]