Schedule resource mutation (#13193)

# Objective

Resolves #13185 

## Solution

Move the following methods from `sub_app` to the `Schedules` resource,
and use them in the sub app:

- `add_systems`
- `configure_sets`
- `ignore_ambiguity`

Add an `entry(&mut self, label: impl ScheduleLabel) -> &mut Schedule`
method to the `Schedules` resource, which returns a mutable reference to
the schedule associated with the label, and creates one if it doesn't
already exist. (build on top of the `entry(..).or_insert_with(...)`
pattern in `HashMap`.

## Testing

- Did you test these changes? If so, how? Added 4 unit tests to the
`schedule.rs` - one that validates adding a system to an existing
schedule, one that validates adding a system to a new one, one that
validates configuring sets on an existing schedule, and one that
validates configuring sets on a new schedule.
- I didn't add tests for `entry` since the previous 4 tests use
functions that rely on it.
- I didn't test `ignore_ambiguity` since I didn't see examples of it's
use, and am not familiar enough with it to know how to set up a good
test for it. However, it relies on the `entry` method as well, so it
should work just like the other 2 methods.

crates/bevy_app/src/sub_app.rs

@@ -177,15 +177,8 @@ impl SubApp {
         schedule: impl ScheduleLabel,
         systems: impl IntoSystemConfigs<M>,
     ) -> &mut Self {
-        let label = schedule.intern();
         let mut schedules = self.world.resource_mut::<Schedules>();
-        if let Some(schedule) = schedules.get_mut(label) {
-            schedule.add_systems(systems);
-        } else {
-            let mut new_schedule = Schedule::new(label);
-            new_schedule.add_systems(systems);
-            schedules.insert(new_schedule);
-        }
+        schedules.add_systems(schedule, systems);
 
         self
     }
@@ -205,15 +198,8 @@ impl SubApp {
         schedule: impl ScheduleLabel,
         sets: impl IntoSystemSetConfigs,
     ) -> &mut Self {
-        let label = schedule.intern();
         let mut schedules = self.world.resource_mut::<Schedules>();
-        if let Some(schedule) = schedules.get_mut(label) {
-            schedule.configure_sets(sets);
-        } else {
-            let mut new_schedule = Schedule::new(label);
-            new_schedule.configure_sets(sets);
-            schedules.insert(new_schedule);
-        }
+        schedules.configure_sets(schedule, sets);
         self
     }
 
@@ -304,14 +290,7 @@ impl SubApp {
         let schedule = schedule.intern();
         let mut schedules = self.world.resource_mut::<Schedules>();
 
-        if let Some(schedule) = schedules.get_mut(schedule) {
-            let schedule: &mut Schedule = schedule;
-            schedule.ignore_ambiguity(a, b);
-        } else {
-            let mut new_schedule = Schedule::new(schedule);
-            new_schedule.ignore_ambiguity(a, b);
-            schedules.insert(new_schedule);
-        }
+        schedules.ignore_ambiguity(schedule, a, b);
 
         self
     }

crates/bevy_ecs/src/schedule/schedule.rs

@@ -78,6 +78,13 @@ impl Schedules {
         self.inner.get_mut(&label.intern())
     }
 
+    /// Returns a mutable reference to the schedules associated with `label`, creating one if it doesn't already exist.
+    pub fn entry(&mut self, label: impl ScheduleLabel) -> &mut Schedule {
+        self.inner
+            .entry(label.intern())
+            .or_insert_with(|| Schedule::new(label))
+    }
+
     /// Returns an iterator over all schedules. Iteration order is undefined.
     pub fn iter(&self) -> impl Iterator<Item = (&dyn ScheduleLabel, &Schedule)> {
         self.inner
@@ -146,6 +153,51 @@ impl Schedules {
 
         info!("{}", message);
     }
+
+    /// Adds one or more systems to the [`Schedule`] matching the provided [`ScheduleLabel`].
+    pub fn add_systems<M>(
+        &mut self,
+        schedule: impl ScheduleLabel,
+        systems: impl IntoSystemConfigs<M>,
+    ) -> &mut Self {
+        self.entry(schedule).add_systems(systems);
+
+        self
+    }
+
+    /// Configures a collection of system sets in the provided schedule, adding any sets that do not exist.
+    #[track_caller]
+    pub fn configure_sets(
+        &mut self,
+        schedule: impl ScheduleLabel,
+        sets: impl IntoSystemSetConfigs,
+    ) -> &mut Self {
+        self.entry(schedule).configure_sets(sets);
+
+        self
+    }
+
+    /// Suppress warnings and errors that would result from systems in these sets having ambiguities
+    /// (conflicting access but indeterminate order) with systems in `set`.
+    ///
+    /// When possible, do this directly in the `.add_systems(Update, a.ambiguous_with(b))` call.
+    /// However, sometimes two independent plugins `A` and `B` are reported as ambiguous, which you
+    /// can only suppress as the consumer of both.
+    #[track_caller]
+    pub fn ignore_ambiguity<M1, M2, S1, S2>(
+        &mut self,
+        schedule: impl ScheduleLabel,
+        a: S1,
+        b: S2,
+    ) -> &mut Self
+    where
+        S1: IntoSystemSet<M1>,
+        S2: IntoSystemSet<M2>,
+    {
+        self.entry(schedule).ignore_ambiguity(a, b);
+
+        self
+    }
 }
 
 fn make_executor(kind: ExecutorKind) -> Box<dyn SystemExecutor> {
@@ -1985,16 +2037,21 @@ pub struct ScheduleNotInitialized;
 
 #[cfg(test)]
 mod tests {
+    use bevy_ecs_macros::ScheduleLabel;
+
     use crate::{
         self as bevy_ecs,
         prelude::{Res, Resource},
         schedule::{
-            IntoSystemConfigs, IntoSystemSetConfigs, Schedule, ScheduleBuildSettings, SystemSet,
+            tests::ResMut, IntoSystemConfigs, IntoSystemSetConfigs, Schedule,
+            ScheduleBuildSettings, SystemSet,
         },
         system::Commands,
         world::World,
     };
 
+    use super::Schedules;
+
     #[derive(Resource)]
     struct Resource1;
 
@@ -2452,4 +2509,127 @@ mod tests {
             });
         }
     }
+
+    #[derive(ScheduleLabel, Hash, Debug, Clone, PartialEq, Eq)]
+    struct TestSchedule;
+
+    #[derive(Resource)]
+    struct CheckSystemRan(usize);
+
+    #[test]
+    fn add_systems_to_existing_schedule() {
+        let mut schedules = Schedules::default();
+        let schedule = Schedule::new(TestSchedule);
+
+        schedules.insert(schedule);
+        schedules.add_systems(TestSchedule, |mut ran: ResMut<CheckSystemRan>| ran.0 += 1);
+
+        let mut world = World::new();
+
+        world.insert_resource(CheckSystemRan(0));
+        world.insert_resource(schedules);
+        world.run_schedule(TestSchedule);
+
+        let value = world
+            .get_resource::<CheckSystemRan>()
+            .expect("CheckSystemRan Resource Should Exist");
+        assert_eq!(value.0, 1);
+    }
+
+    #[test]
+    fn add_systems_to_non_existing_schedule() {
+        let mut schedules = Schedules::default();
+
+        schedules.add_systems(TestSchedule, |mut ran: ResMut<CheckSystemRan>| ran.0 += 1);
+
+        let mut world = World::new();
+
+        world.insert_resource(CheckSystemRan(0));
+        world.insert_resource(schedules);
+        world.run_schedule(TestSchedule);
+
+        let value = world
+            .get_resource::<CheckSystemRan>()
+            .expect("CheckSystemRan Resource Should Exist");
+        assert_eq!(value.0, 1);
+    }
+
+    #[derive(SystemSet, Debug, Hash, Clone, PartialEq, Eq)]
+    enum TestSet {
+        First,
+        Second,
+    }
+
+    #[test]
+    fn configure_set_on_existing_schedule() {
+        let mut schedules = Schedules::default();
+        let schedule = Schedule::new(TestSchedule);
+
+        schedules.insert(schedule);
+
+        schedules.configure_sets(TestSchedule, (TestSet::First, TestSet::Second).chain());
+        schedules.add_systems(
+            TestSchedule,
+            (|mut ran: ResMut<CheckSystemRan>| {
+                assert_eq!(ran.0, 0);
+                ran.0 += 1;
+            })
+            .in_set(TestSet::First),
+        );
+
+        schedules.add_systems(
+            TestSchedule,
+            (|mut ran: ResMut<CheckSystemRan>| {
+                assert_eq!(ran.0, 1);
+                ran.0 += 1;
+            })
+            .in_set(TestSet::Second),
+        );
+
+        let mut world = World::new();
+
+        world.insert_resource(CheckSystemRan(0));
+        world.insert_resource(schedules);
+        world.run_schedule(TestSchedule);
+
+        let value = world
+            .get_resource::<CheckSystemRan>()
+            .expect("CheckSystemRan Resource Should Exist");
+        assert_eq!(value.0, 2);
+    }
+
+    #[test]
+    fn configure_set_on_new_schedule() {
+        let mut schedules = Schedules::default();
+
+        schedules.configure_sets(TestSchedule, (TestSet::First, TestSet::Second).chain());
+        schedules.add_systems(
+            TestSchedule,
+            (|mut ran: ResMut<CheckSystemRan>| {
+                assert_eq!(ran.0, 0);
+                ran.0 += 1;
+            })
+            .in_set(TestSet::First),
+        );
+
+        schedules.add_systems(
+            TestSchedule,
+            (|mut ran: ResMut<CheckSystemRan>| {
+                assert_eq!(ran.0, 1);
+                ran.0 += 1;
+            })
+            .in_set(TestSet::Second),
+        );
+
+        let mut world = World::new();
+
+        world.insert_resource(CheckSystemRan(0));
+        world.insert_resource(schedules);
+        world.run_schedule(TestSchedule);
+
+        let value = world
+            .get_resource::<CheckSystemRan>()
+            .expect("CheckSystemRan Resource Should Exist");
+        assert_eq!(value.0, 2);
+    }
 }

crates/bevy_ecs/src/schedule/state.rs

@@ -375,22 +375,9 @@ pub fn setup_state_transitions_in_world(
     schedules.insert(schedule);
 
     if let Some(startup) = startup_label {
-        match schedules.get_mut(startup) {
-            Some(schedule) => {
-                schedule.add_systems(|world: &mut World| {
-                    let _ = world.try_run_schedule(StateTransition);
-                });
-            }
-            None => {
-                let mut schedule = Schedule::new(startup);
-
-                schedule.add_systems(|world: &mut World| {
-                    let _ = world.try_run_schedule(StateTransition);
-                });
-
-                schedules.insert(schedule);
-            }
-        }
+        schedules.add_systems(startup, |world: &mut World| {
+            let _ = world.try_run_schedule(StateTransition);
+        });
     }
 }