Run fixed time-step in an exclusive system

crates/bevy_app/Cargo.toml

@@ -25,6 +25,7 @@ bevy_tasks = { path = "../bevy_tasks", version = "0.8.0-dev" }
 # other
 serde = { version = "1.0", features = ["derive"], optional = true }
 ron = { version = "0.7.0", optional = true }
+thiserror = "1.0"
 
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]

crates/bevy_app/src/app.rs

@@ -1,11 +1,15 @@
-use crate::{CoreStage, Plugin, PluginGroup, PluginGroupBuilder, StartupSchedule, StartupStage};
+use crate::{
+    CoreStage, IntoSubSchedule, Plugin, PluginGroup, PluginGroupBuilder, StartupSchedule,
+    StartupStage, SubSchedules,
+};
 pub use bevy_derive::AppLabel;
 use bevy_ecs::{
+    change_detection::Mut,
     event::{Event, Events},
     prelude::{FromWorld, IntoExclusiveSystem},
     schedule::{
-        IntoSystemDescriptor, Schedule, ShouldRun, Stage, StageLabel, State, StateData, SystemSet,
-        SystemStage,
+        IntoSystemDescriptor, Schedule, ScheduleLabel, ShouldRun, Stage, StageLabel, State,
+        StateData, SystemSet, SystemStage,
     },
     system::Resource,
     world::World,
@@ -311,6 +315,87 @@ impl App {
         self
     }
 
+    /// Adds a [`SubSchedule`](crate::SubSchedule) to the app.
+    ///
+    /// If a [label](bevy_ecs::schedule::ScheduleLabel) has been applied to it,
+    /// it can be accessed in the resource [`SubSchedules`].
+    ///
+    /// # Examples
+    /// ```
+    /// # use bevy_app::prelude::*;
+    /// # use bevy_ecs::prelude::*;
+    /// # let mut app = App::new();
+    /// # fn foo_system() {}
+    /// # fn bar_system() {}
+    /// #
+    /// // Define a label for the a schedule.
+    /// #[derive(ScheduleLabel)]
+    /// struct MySched;
+    ///
+    /// // Add a new sub-schedule. If the schedule only needs one stage,
+    /// // you can avoid some overhead by adding the stage directly.
+    /// app.add_sub_schedule(
+    ///     SystemStage::parallel()
+    ///         .with_system(foo_system)
+    ///         .with_system(bar_system)
+    ///         .label(MySched),
+    /// );
+    /// ```
+    #[track_caller]
+    pub fn add_sub_schedule(&mut self, sched: impl IntoSubSchedule) -> &mut Self {
+        crate::sub_schedule::add_to_app(self, IntoSubSchedule::into_sched(sched));
+        self
+    }
+
+    /// Allows you to modify a [`SubSchedule`](crate::SubSchedule) by executing a closure on it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use bevy_app::prelude::*;
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # #[derive(ScheduleLabel)]
+    /// # struct FixedUpdate;
+    /// #
+    /// # let mut app = App::new();
+    /// # app.add_sub_schedule(
+    /// #   Schedule::default()
+    /// #       .with_stage(CoreStage::Update, SystemStage::parallel())
+    /// #       .label(FixedUpdate),
+    /// # );
+    /// # fn my_system() {}
+    /// #
+    /// // Access a previously-added sub-schedule and add a system to it.
+    /// app.sub_schedule(FixedUpdate, |sched: &mut Schedule| {
+    ///     sched.add_system_to_stage(CoreStage::Update, my_system);
+    /// });
+    /// ```
+    ///
+    /// # Panics
+    /// If `label` refers to a non-existent `SubSchedule`, or if it is not of type `S`.
+    pub fn sub_schedule<S: Stage, F>(&mut self, label: impl ScheduleLabel, f: F) -> &mut Self
+    where
+        F: FnOnce(&mut S),
+    {
+        #[inline(never)]
+        fn panic(label: impl Debug, ty: impl Debug) -> ! {
+            panic!("there is no sub-schedule labeled '{label:?}', or it does not match the type `{ty:?}`");
+        }
+
+        let label = label.as_label();
+
+        // Make sure the resource is initialized, so we get a more helpful panic message.
+        let mut schedules: Mut<SubSchedules> =
+            self.world.get_resource_or_insert_with(Default::default);
+        let sched: &mut S = match schedules.get_mut(label) {
+            Some(x) => x,
+            None => panic(label, std::any::type_name::<S>()),
+        };
+        f(sched);
+        self
+    }
+
     /// Adds a system to the [update stage](Self::add_default_stages) of the app's [`Schedule`].
     ///
     /// Refer to the [system module documentation](bevy_ecs::system) to see how a system

crates/bevy_app/src/lib.rs

@@ -6,6 +6,7 @@ mod app;
 mod plugin;
 mod plugin_group;
 mod schedule_runner;
+mod sub_schedule;
 
 #[cfg(feature = "bevy_ci_testing")]
 mod ci_testing;
@@ -15,12 +16,14 @@ pub use bevy_derive::DynamicPlugin;
 pub use plugin::*;
 pub use plugin_group::*;
 pub use schedule_runner::*;
+pub use sub_schedule::*;
 
 #[allow(missing_docs)]
 pub mod prelude {
     #[doc(hidden)]
     pub use crate::{
-        app::App, CoreStage, DynamicPlugin, Plugin, PluginGroup, StartupSchedule, StartupStage,
+        app::App, CoreStage, DynamicPlugin, IntoSubSchedule, Plugin, PluginGroup, StartupSchedule,
+        StartupStage,
     };
 }
 

crates/bevy_app/src/sub_schedule.rs

@@ -0,0 +1,239 @@
+use std::fmt::{Debug, Display};
+
+use bevy_ecs::{
+    change_detection::Mut,
+    schedule::{ScheduleLabel, ScheduleLabelId, Stage, StageLabel, StageLabelId},
+    system::IntoExclusiveSystem,
+    world::World,
+};
+use bevy_utils::HashMap;
+use thiserror::Error;
+
+use crate::App;
+
+/// Methods for converting a schedule into a [sub-Schedule](SubSchedule) descriptor.
+pub trait IntoSubSchedule: Sized {
+    /// The type that controls the behaviour of the exclusive system
+    /// which runs [`SubSchedule`]s of this type.
+    type Runner: FnMut(&mut dyn Stage, &mut World) + Send + Sync + 'static;
+
+    /// Applies the specified label to the current schedule.
+    /// This means it will be accessible in the [`SubSchedules`] resource after
+    /// being added to the [`App`].
+    fn label(self, label: impl ScheduleLabel) -> SubSchedule<Self::Runner> {
+        let mut sub = Self::into_sched(self);
+        sub.label = Some(label.as_label());
+        sub
+    }
+    /// Defines a function that runs the current schedule. It will be inserted into
+    /// an exclusive system within the stage `stage`.
+    ///
+    /// Overwrites any previously set runner or stage.
+    fn with_runner<F>(self, stage: impl StageLabel, f: F) -> SubSchedule<F>
+    where
+        F: FnMut(&mut dyn Stage, &mut World) + Send + Sync + 'static,
+    {
+        let SubSchedule {
+            schedule, label, ..
+        } = Self::into_sched(self);
+        SubSchedule {
+            schedule,
+            label,
+            runner: Some((stage.as_label(), f)),
+        }
+    }
+
+    /// Performs the conversion. You usually do not need to call this directly.
+    fn into_sched(_: Self) -> SubSchedule<Self::Runner>;
+}
+
+impl<S: Stage> IntoSubSchedule for S {
+    type Runner = fn(&mut dyn Stage, &mut World);
+    fn into_sched(schedule: Self) -> SubSchedule<Self::Runner> {
+        SubSchedule {
+            schedule: Box::new(schedule),
+            label: None,
+            runner: None,
+        }
+    }
+}
+
+impl<R> IntoSubSchedule for SubSchedule<R>
+where
+    R: FnMut(&mut dyn Stage, &mut World) + Send + Sync + 'static,
+{
+    type Runner = R;
+    #[inline]
+    fn into_sched(sched: Self) -> SubSchedule<R> {
+        sched
+    }
+}
+
+/// A schedule that may run independently of the main app schedule.
+pub struct SubSchedule<F>
+where
+    F: FnMut(&mut dyn Stage, &mut World) + Send + Sync + 'static,
+{
+    schedule: Box<dyn Stage>,
+    label: Option<ScheduleLabelId>,
+    runner: Option<(StageLabelId, F)>,
+}
+
+/// A [resource](bevy_ecs::system::Res) that stores all labeled [`SubSchedule`]s.
+#[derive(Default)]
+pub struct SubSchedules {
+    // INVARIANT: A `SubSlot` cannot be removed once added, and is associated with
+    // a single schedule. Even if a slot gets temporarily emptied, it is guaranteed
+    // that the slot will always get refilled by the same exact schedule.
+    map: HashMap<ScheduleLabelId, SubSlot>,
+}
+
+struct SubSlot(Option<Box<dyn Stage>>);
+
+/// Error type returned by [`SubSchedules::extract_from`](SubSchedules#method.extract_from).
+#[derive(Error)]
+#[non_exhaustive]
+pub enum ExtractError {
+    /// Schedule could not be found.
+    #[error("there is no sub-schedule with label '{0:?}'")]
+    NotFound(ScheduleLabelId),
+    /// Schedule is being extracted by someone else right now.
+    #[error("cannot extract sub-schedule '{0:?}', as it is currently extracted already")]
+    AlreadyExtracted(ScheduleLabelId),
+    /// The [`SubSchedules`] resource got removed during the scope.
+    #[error("the `SubSchedules` resource got removed during the scope for schedule '{0:?}'")]
+    ResourceLost(ScheduleLabelId),
+}
+impl Debug for ExtractError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        Display::fmt(self, f)
+    }
+}
+
+#[derive(Error)]
+/// Error type returned by [`SubSchedules::insert`](SubSchedules#method.insert).
+#[non_exhaustive]
+pub enum InsertError {
+    /// A schedule with this label already exists.
+    #[error("a sub-schedule with label '{0:?}' already exists")]
+    Duplicate(ScheduleLabelId),
+}
+impl Debug for InsertError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        Display::fmt(self, f)
+    }
+}
+
+impl SubSchedules {
+    /// Inserts a new sub-schedule.
+    ///
+    /// # Errors
+    /// If there is already a sub-schedule labeled `label`.
+    pub fn insert(
+        &mut self,
+        label: impl ScheduleLabel,
+        sched: Box<dyn Stage>,
+    ) -> Result<(), InsertError> {
+        let label = label.as_label();
+        if self.map.contains_key(&label) {
+            return Err(InsertError::Duplicate(label));
+        }
+        self.map.insert(label, SubSlot(Some(sched)));
+        Ok(())
+    }
+
+    /// Temporarily extracts a [`SubSchedule`] from the world, and provides a scope
+    /// that has mutable access to both the schedule and the [`World`].
+    /// At the end of this scope, the sub-schedule is automatically reinserted.
+    ///
+    /// # Errors
+    /// If there is no schedule associated with `label`, or if that schedule
+    /// is currently already extracted.
+    pub fn extract_scope<F, T>(
+        world: &mut World,
+        label: impl ScheduleLabel,
+        f: F,
+    ) -> Result<T, ExtractError>
+    where
+        F: FnOnce(&mut World, &mut dyn Stage) -> T,
+    {
+        let label = label.as_label();
+
+        // Extract.
+        let mut schedules = world.resource_mut::<Self>();
+        let mut sched = schedules
+            .map
+            .get_mut(&label)
+            .ok_or(ExtractError::NotFound(label))?
+            .0
+            .take()
+            .ok_or(ExtractError::AlreadyExtracted(label))?;
+
+        // Execute.
+        let val = f(world, sched.as_mut());
+
+        // Re-insert.
+        world
+            .get_resource_mut::<Self>()
+            .ok_or(ExtractError::ResourceLost(label))?
+            .map
+            .get_mut(&label)
+            .expect("schedule must exist in the map since we found it previously")
+            .0 = Some(sched);
+
+        Ok(val)
+    }
+
+    /// Gets a mutable reference to the sub-schedule identified by `label`.
+    ///
+    /// # Panics
+    /// If the schedule is currently [extracted](#method.extract_scope).
+    pub fn get_mut<S: Stage>(&mut self, label: impl ScheduleLabel) -> Option<&mut S> {
+        #[cold]
+        fn panic(label: impl Debug) -> ! {
+            panic!("cannot get sub-schedule '{label:?}', as it is currently extracted")
+        }
+
+        let label = label.as_label();
+        let sched = match self.map.get_mut(&label)?.0.as_deref_mut() {
+            Some(x) => x,
+            None => panic(label),
+        };
+        sched.downcast_mut()
+    }
+}
+
+#[track_caller]
+pub(crate) fn add_to_app(app: &mut App, schedule: impl IntoSubSchedule) {
+    let SubSchedule {
+        mut schedule,
+        label,
+        runner,
+    } = IntoSubSchedule::into_sched(schedule);
+
+    // If it has a label, insert it to the public resource.
+    if let Some(label) = label {
+        let mut res: Mut<SubSchedules> = app.world.get_resource_or_insert_with(Default::default);
+        res.insert(label, schedule).unwrap();
+
+        if let Some((stage, mut runner)) = runner {
+            // Driver which extracts the schedule from the world and runs it.
+            let driver = move |w: &mut World| {
+                SubSchedules::extract_scope(w, label, |w, sched| {
+                    runner(sched, w);
+                })
+                .unwrap();
+            };
+            app.add_system_to_stage(stage, driver.exclusive_system());
+        }
+    } else if let Some((stage, mut runner)) = runner {
+        // If there's no label, then the schedule isn't visible publicly.
+        // We can just store it locally
+        let driver = move |w: &mut World| {
+            runner(schedule.as_mut(), w);
+        };
+        app.add_system_to_stage(stage, driver.exclusive_system());
+    } else {
+        panic!("inserted sub-schedule can never be accessed, as it has neither a label nor a runner function")
+    }
+}