Bevy Staged Plugin Example/Template

main.rs

use bevy::prelude::*;

use self::{
    labels::{StageCleanup, StageSet, StageSetup},
    resources::{StageConfigs, StageSetList, SystemConfigsFactory},
};

mod labels {
    use bevy::prelude::*;

    /// An anonymous [`SystemSet`] used for a single iteration of the systems added via
    /// [`super::StagedPlugin::add_systems`].
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub(crate) struct StageSet(usize);

    impl StageSet {
        pub(crate) fn new() -> Self {
            Self(Self::next_id())
        }

        /// Gets the numeric id of this [`StageSet`]. Note that this id is only used to provide uniqueness while
        /// hashing, and is not used otherwise.
        pub fn id(self) -> usize {
            self.0
        }

        /// Atomically produces an increasing ID
        fn next_id() -> usize {
            static NEXT_ID: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(1);
            NEXT_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
        }
    }

    impl SystemSet for StageSet {
        fn dyn_clone(&self) -> Box<dyn SystemSet> {
            Box::new(*self)
        }
    }

    /// A [`SystemSet`] that runs before all [`StageSet`]s.
    #[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct StageSetup;

    /// A [`SystemSet`] that runs after all [`StageSet`]s.
    #[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct StageCleanup;
}

mod resources {
    use bevy::{ecs::schedule::SystemConfigs, prelude::*};

    use super::labels::StageSet;

    /// Contains the list of anonymous [`StageSet`]s, in order. Produced by [`super::StagedPlugin`], and removed during
    /// startup.
    #[derive(Resource)]
    pub(crate) struct StageSetList(pub(crate) Vec<StageSet>);

    /// Contains the set of factory functions that provide system configs to be added to every [`StageSet`]. Created
    /// implicitly by any call to [`super::StagedPlugin::add_systems`], and removed during startup.
    #[derive(Default)]
    pub(crate) struct StageConfigs(pub(crate) Vec<SystemConfigsFactory>);

    /// A helper type for functions that produce [`SystemConfigs`].
    pub(crate) type SystemConfigsFactory = &'static dyn Fn() -> SystemConfigs;
}

/// A [`Plugin`] demonstrating how to set up a looped [`SystemSet`] without needing an exclusive system or an extra
/// schedule.
///
/// This core plugin is configured with the number of iterations to perform, and creates a list of anonymous sets (see:
/// [`labels::StageSet`]) that are added to the [`Update`] schedule and have their relative ordering enforced normally.
/// Then this plugin, or other plugins, can use [`StagedPlugin::add_systems`] to add systems to *every* stage. This is
/// accomplished by requiring `add_systems` to accept a factory function that produces a [`SystemConfigs`], storing
/// that factory in a resource, then using the factory to add the systems to each set during [`Startup`].
/// 
/// [`SystemConfigs`]: bevy::ecs::schedule::SystemConfigs
pub struct StagedPlugin {
    pub num_stages: usize,
}

impl StagedPlugin {
    /// Adds the given system(s) to the update stage.
    ///
    /// Note: A systems configuration can be converted to [`SystemConfigs`] using the `into_configs` method.
    /// ToDo: Find a way to accept `impl IntoSystemConfigs` so this function is more similar to [`App::add_systems`].
    ///
    /// ```ignore
    /// add_systems(app, &|| my_system.into_configs());
    /// ```
    pub fn add_systems(app: &mut App, systems: SystemConfigsFactory) {
        app.world.init_non_send_resource::<StageConfigs>();
        app.world
            .get_non_send_resource_mut::<StageConfigs>()
            .unwrap()
            .0
            .push(systems);
    }
}

impl Plugin for StagedPlugin {
    fn build(&self, app: &mut App) {
        // Create an anonymous set for each stage
        let stage_sets = (0..self.num_stages)
            .map(|_| StageSet::new())
            .collect::<Vec<_>>();

        // Ensure stage ordering:
        // StageSetup -> Stage0 -> Stage1 -> ... -> StageN -> StageCleanup
        if stage_sets.len() > 0 {
            let first_stage = stage_sets[0];
            let last_stage = stage_sets[stage_sets.len() - 1];

            app.configure_set(Update, StageSetup.before(first_stage));

            for i in 1..stage_sets.len() {
                app.configure_set(Update, stage_sets[i].after(stage_sets[i - 1]));
            }

            app.configure_set(Update, StageCleanup.after(last_stage));
        } else {
            // If there are 0 update stages, just put cleanup directly after setup
            app.configure_set(Update, StageSetup.before(StageCleanup));
        }

        // Add the list of stage sets to the world as a resource. This will be consumed by `init_staged_plugin_system`
        // on startup so that systems can be added to every stage.
        app.insert_resource(StageSetList(stage_sets))
            .add_systems(Startup, init_staged_plugin_system);

        // Extra logging systems - not necessary for the plugin to function
        app.add_systems(
            Update,
            (
                setup_system.in_set(StageSetup),
                cleanup_system.in_set(StageCleanup),
            ),
        );
    }
}

/// Consumes the [`StageSetList`] and [`StageConfigs`] resources, and adds all configs to **every** [`StageSet`].
fn init_staged_plugin_system(world: &mut World) {
    // Get resources
    let stage_sets = world.remove_resource::<StageSetList>().unwrap().0;
    let stage_configs = world
        .remove_non_send_resource::<StageConfigs>()
        .unwrap_or_default()
        .0;

    // Get the `Update` schedule
    let mut schedules = world.resource_mut::<Schedules>();
    let schedule = schedules
        .get_mut(&Update)
        .expect("Update schedule does not exist");

    // For each stage set...
    for set in stage_sets.into_iter() {
        // Extra logging systems - not necessary for the plugin to function
        schedule.add_systems((move || println!("Running set {}", set.id())).in_set(set));

        // And each stage system...
        for config in stage_configs.iter().copied() {
            // Create a config for the system, and add it to the set
            schedule.add_systems(config().in_set(set));
        }
    }
}

fn setup_system() {
    println!("setup");
}

fn cleanup_system() {
    println!("cleanup");
}

/// A simple plugin that shows how other plugins can make use of [`StagedPlugin`]'s stages.
pub struct StageUserPlugin;

impl Plugin for StageUserPlugin {
    fn build(&self, app: &mut App) {
        StagedPlugin::add_systems(app, &|| stage_system.into_configs());
    }
}

fn stage_system() {
    println!("stage user system");
}

fn main() {
    let mut app = App::new();

    app.add_plugins((StagedPlugin { num_stages: 5 }, StageUserPlugin));

    app.world.run_schedule(Main);
}

Possible output:

setup
stage user system
Running set 1
stage user system
Running set 2
Running set 3
stage user system
Running set 4
stage user system
stage user system
Running set 5
cleanup

Note that each "Running set N" runs simultaneously with a "stage user system", so their pairwise ordering is arbitrary. Taken as pairs, though, it can be seen that each set runs in order.