using Schedule labels with custom app states

main.rs

use bevy::prelude::*;
use states::AppUpdate;

mod prelude;
mod states;
use prelude::*;

fn in_menu(mouse: Res<Input<MouseButton>>, mut next: ResMut<NextState<states::AppState>>) {
    println!("in_menu");

    if mouse.just_pressed(MouseButton::Left) {
        println!("left click");

        next.0 = Some(AppState::FileEditing);
    }
}

fn in_edit(mouse: Res<Input<MouseButton>>, mut next: ResMut<NextState<states::AppState>>) {
    println!("in_edit");

    if mouse.just_pressed(MouseButton::Left) {
        println!("left click");

        next.0 = Some(AppState::Message("hello".into()));
    }
}

fn in_message(st: Res<State<states::AppState>>) {
    let AppState::Message(st) = st.get() else {
        unreachable!("in_message called when not in message state")
    };

    println!("message {:?}", st);
}

fn show_state(cur: Res<State<states::AppState>>) {
    println!("state: {:?}", cur.get());
}

fn main() {
    App::new()
        .add_plugins((DefaultPlugins, states::StatesPlugin))
        .add_systems(Update, show_state)
        .add_systems(AppUpdate(MainMenu), in_menu)
        .add_systems(AppUpdate(FileEditing), in_edit)
        .add_systems(AppUpdate(Message), in_message)
        .run();
}

states.rs

use crate::prelude::*;
use bevy::{ecs::schedule::ScheduleLabel, prelude::*, utils::CowArc};
use strum::IntoEnumIterator;
use strum::{EnumDiscriminants, EnumIter};

#[derive(States, Debug, Clone, Eq, PartialEq, Hash, Default, EnumDiscriminants)]
#[strum_discriminants(derive(Hash, EnumIter))]
pub enum AppState {
    #[default]
    MainMenu,
    FileEditing,
    Message(Str),
}

#[derive(ScheduleLabel, Hash, Debug, PartialEq, Eq, Clone)]
pub struct AppUpdate(pub AppStateDiscriminants);

pub struct StatesPlugin;

impl Plugin for StatesPlugin {
    fn build(&self, app: &mut App) {
        app.add_state::<AppState>();

        for discr in AppStateDiscriminants::iter() {
            app.init_schedule(AppUpdate(discr)).add_systems(
                Update,
                (move |world: &mut World| world.try_run_schedule(AppUpdate(discr)).unwrap())
                    .run_if(IntoSystem::into_system(move |a: Res<State<AppState>>| {
                        discr == a.get().into()
                    })),
            );
        }
    }
}