feat: add support for auto_init_state and auto_register_state_type

Cargo.toml

@@ -30,6 +30,7 @@ bevy_auto_plugin_nightly_shared = { version = "0.1.0", path = "crates/bevy_auto_
 bevy_auto_plugin_proc_macros = { version = "0.1.0", path = "crates/bevy_auto_plugin_proc_macros" }
 bevy_app = { version = "0.15", default-features = false }
 bevy_core = { version = "0.15" }
+bevy_state = { version = "0.15" }
 bevy_reflect = { version = "0.15" }
 bevy_internal = { version = "0.15" }
 bevy_ecs = { version = "0.15" }
@@ -48,6 +49,7 @@ bevy_app = { workspace = true, default-features = false }
 bevy_core = { workspace = true }
 bevy_reflect = { workspace = true }
 bevy_ecs = { workspace = true }
+bevy_state = { workspace = true }
 internal_test_util = { path = "crates/internal_test_util" }
 trybuild = "1.0"
 log = { workspace = true }

crates/bevy_auto_plugin_nightly_proc_macros/src/lib.rs

@@ -13,8 +13,8 @@ use bevy_auto_plugin_shared::util::{
     resolve_path_from_item_or_args, FnParamMutabilityCheckErrMessages, Target,
 };
 use bevy_auto_plugin_shared::{
-    generate_add_events, generate_auto_names, generate_init_resources, generate_register_types,
-    util,
+    generate_add_events, generate_auto_names, generate_init_resources, generate_init_states,
+    generate_register_state_types, generate_register_types, util,
 };
 use proc_macro2::{Ident, Span};
 use quote::quote;
@@ -168,16 +168,24 @@ fn auto_plugin_inner(file_path: String, app_param_name: &Ident) -> Result<MacroS
             app_param_name,
             file_state.context.register_types.clone().drain(),
         )?;
+        let register_state_types = generate_register_state_types(
+            app_param_name,
+            file_state.context.register_state_types.drain(),
+        )?;
         let add_events =
             generate_add_events(app_param_name, file_state.context.add_events.drain())?;
         let init_resources =
             generate_init_resources(app_param_name, file_state.context.init_resources.drain())?;
+        let init_states =
+            generate_init_states(app_param_name, file_state.context.init_states.drain())?;
         let auto_names =
             generate_auto_names(app_param_name, file_state.context.auto_names.drain())?;
         Ok(quote! {
             #register_types
+            #register_state_types
             #add_events
             #init_resources
+            #init_states
             #auto_names
         })
     })
@@ -373,3 +381,48 @@ pub fn auto_init_resource(attr: CompilerStream, input: CompilerStream) -> Compil
 pub fn auto_name(attr: CompilerStream, input: CompilerStream) -> CompilerStream {
     handle_attribute(attr, input, Target::RequiredComponentAutoName)
 }
+
+/// Automatically initializes a State in the Bevy `App`.
+///
+/// # Example (without generics)
+/// ```no_run
+/// use bevy::prelude::*;
+/// use bevy_auto_plugin::auto_plugin::*;
+///
+/// #[auto_init_state]
+/// #[derive(States, Debug, Copy, Clone, Default, PartialEq, Eq, Hash)]
+/// struct Foo;
+///
+/// #[auto_plugin(app=app)]
+/// fn plugin(app: &mut App) {
+///     // generated code:
+///     app.init_state::<Foo>();
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn auto_init_state(attr: CompilerStream, input: CompilerStream) -> CompilerStream {
+    handle_attribute(attr, input, Target::InitStates)
+}
+
+/// Automatically registers a State type in the Bevy `App`.
+///
+/// # Example (without generics)
+/// ```no_run
+/// use bevy::prelude::*;
+/// use bevy_auto_plugin::auto_plugin::*;
+///
+/// #[auto_register_state_type]
+/// #[derive(States, Debug, Copy, Clone, Default, PartialEq, Eq, Hash, Reflect)]
+/// struct Foo;
+///
+/// #[auto_plugin(app=app)]
+/// fn plugin(app: &mut App) {
+///     // generated code:
+///     app.register_type::<State<Foo>>();
+///     app.register_type::<NextState<Foo>>();
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn auto_register_state_type(attr: CompilerStream, input: CompilerStream) -> CompilerStream {
+    handle_attribute(attr, input, Target::RegisterStateTypes)
+}

crates/bevy_auto_plugin_nightly_shared/src/lib.rs

@@ -58,8 +58,10 @@ pub fn update_state(
         let path = path_to_string(&path, false);
         let inserted = match target {
             Target::RegisterTypes => entry.context.register_types.insert(path),
+            Target::RegisterStateTypes => entry.context.register_state_types.insert(path),
             Target::AddEvents => entry.context.add_events.insert(path),
             Target::InitResources => entry.context.init_resources.insert(path),
+            Target::InitStates => entry.context.init_states.insert(path),
             Target::RequiredComponentAutoName => entry.context.auto_names.insert(path),
         };
         if !inserted {

crates/bevy_auto_plugin_proc_macros/src/lib.rs

@@ -5,7 +5,8 @@ use bevy_auto_plugin_shared::util::{
     inject_module, items_with_attribute_macro, ItemWithAttributeMatch,
 };
 use bevy_auto_plugin_shared::{
-    generate_add_events, generate_auto_names, generate_init_resources, generate_register_types,
+    generate_add_events, generate_auto_names, generate_init_resources, generate_init_states,
+    generate_register_state_types, generate_register_types,
 };
 use proc_macro2::{Ident, Span};
 use quote::quote;
@@ -100,18 +101,30 @@ fn auto_plugin_inner(mut module: ItemMod, init_name: &Ident) -> Result<MacroStre
         let auto_names = items_with_attribute_macro(items, "auto_name")?;
         let auto_names = map_to_string(auto_names);
 
+        let auto_register_state_types =
+            items_with_attribute_macro(items, "auto_register_state_type")?;
+        let auto_register_state_types = map_to_string(auto_register_state_types);
+
+        let auto_init_states = items_with_attribute_macro(items, "auto_init_state")?;
+        let auto_init_states = map_to_string(auto_init_states);
+
         inject_module(&mut module, move || {
             let auto_register_types =
                 generate_register_types(&app_param_ident, auto_register_types)?;
             let auto_add_events = generate_add_events(&app_param_ident, auto_add_events)?;
             let auto_init_resources =
                 generate_init_resources(&app_param_ident, auto_init_resources)?;
             let auto_names = generate_auto_names(&app_param_ident, auto_names)?;
+            let auto_register_state_types =
+                generate_register_state_types(&app_param_ident, auto_register_state_types)?;
+            let auto_init_states = generate_init_states(&app_param_ident, auto_init_states)?;
             parse2::<Item>(quote! {
                 pub(super) fn #init_name(app: &mut bevy_app::prelude::App) {
                     #auto_register_types
+                    #auto_register_state_types
                     #auto_add_events
                     #auto_init_resources
+                    #auto_init_states
                     #auto_names
                 }
             })
@@ -353,3 +366,72 @@ pub fn auto_name(_attr: CompilerStream, input: CompilerStream) -> CompilerStream
     // Just return the input unchanged; this acts as a marker.
     input
 }
+
+/// Automatically initializes a State in the Bevy `App`.
+///
+/// # Example (without generics)
+/// ```no_run
+/// use bevy::prelude::*;
+/// use bevy_auto_plugin::auto_plugin::*;
+///
+/// #[auto_plugin(init_name=init)]
+/// pub mod my_plugin {
+///     use super::*;
+///
+///     #[auto_init_state]
+///     #[derive(States, Debug, Copy, Clone, Default, PartialEq, Eq, Hash)]
+///     enum Foo {
+///         #[default]
+///         A,
+///     }
+///
+///     // code gen:
+///     pub(super) fn init(app: &mut App) {  
+///         app.init_state::<FooResource>();
+///     }
+/// }
+///
+/// fn plugin(app: &mut App) {
+///     app.add_plugin(my_plugin::init)
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn auto_init_state(_attr: CompilerStream, input: CompilerStream) -> CompilerStream {
+    // Just return the input unchanged; this acts as a marker.
+    input
+}
+
+/// Automatically registers a State<T> and NextState<T> in the Bevy `App`.
+///
+/// # Example (without generics)
+/// ```no_run
+/// use bevy::prelude::*;
+/// use bevy_auto_plugin::auto_plugin::*;
+///
+/// #[auto_plugin(init_name=init)]
+/// pub mod my_plugin {
+///     use super::*;
+///
+///     #[auto_register_state_type]
+///     #[derive(States, Debug, Copy, Clone, Default, PartialEq, Eq, Hash)]
+///     enum Foo {
+///         #[default]
+///         A,
+///     }
+///
+///     // code gen:
+///     pub(super) fn init(app: &mut App) {  
+///         app.register_type::<State<Foo>>();
+///         app.register_type::<NextState<Foo>>();
+///     }
+/// }
+///
+/// fn plugin(app: &mut App) {
+///     app.add_plugin(my_plugin::init)
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn auto_register_state_type(_attr: CompilerStream, input: CompilerStream) -> CompilerStream {
+    // Just return the input unchanged; this acts as a marker.
+    input
+}

crates/bevy_auto_plugin_shared/src/lib.rs

@@ -9,8 +9,10 @@ pub mod util;
 #[derive(Default)]
 pub struct AutoPluginContext {
     pub register_types: HashSet<String>,
+    pub register_state_types: HashSet<String>,
     pub add_events: HashSet<String>,
     pub init_resources: HashSet<String>,
+    pub init_states: HashSet<String>,
     pub auto_names: HashSet<String>,
 }
 
@@ -94,3 +96,44 @@ pub fn generate_auto_names(
         }
     })
 }
+
+pub fn generate_register_state_types(
+    app_ident: &Ident,
+    items: impl Iterator<Item = String>,
+) -> syn::Result<MacroStream> {
+    let register_state_types = items
+        .map(|item| {
+            let item = syn::parse_str::<Path>(&item)?;
+            Ok(quote! {
+                #app_ident.register_type::<State<#item>>();
+                #app_ident.register_type::<NextState<#item>>();
+            })
+        })
+        .collect::<syn::Result<Vec<_>>>()?;
+    Ok(quote! {
+        {
+            // register_state_types
+            #(#register_state_types)*
+        }
+    })
+}
+
+pub fn generate_init_states(
+    app_ident: &Ident,
+    items: impl Iterator<Item = String>,
+) -> syn::Result<MacroStream> {
+    let init_states = items
+        .map(|item| {
+            let item = syn::parse_str::<Path>(&item)?;
+            Ok(quote! {
+                #app_ident.init_state::<#item>();
+            })
+        })
+        .collect::<syn::Result<Vec<_>>>()?;
+    Ok(quote! {
+        {
+            // init_states
+            #(#init_states)*
+        }
+    })
+}

crates/bevy_auto_plugin_shared/src/util.rs

@@ -53,8 +53,10 @@ pub fn path_to_string(path: &Path, strip_spaces: bool) -> String {
 
 pub enum Target {
     RegisterTypes,
+    RegisterStateTypes,
     AddEvents,
     InitResources,
+    InitStates,
     RequiredComponentAutoName,
 }
 

tests/nightly/auto_init_state.rs

@@ -0,0 +1,52 @@
+use bevy_app::prelude::*;
+use bevy_auto_plugin::auto_plugin::*;
+use bevy_state::app::StatesPlugin;
+use bevy_state::prelude::*;
+
+#[auto_init_state]
+#[derive(States, Debug, Copy, Clone, Default, PartialEq, Eq, Hash)]
+enum Test {
+    #[default]
+    A,
+    #[allow(dead_code)]
+    B,
+}
+
+#[auto_init_state]
+#[derive(SubStates, Debug, Copy, Clone, Default, PartialEq, Eq, Hash)]
+#[source(Test = Test::B)]
+enum InnerTest {
+    #[default]
+    A,
+}
+
+#[auto_plugin(app=app)]
+fn plugin(app: &mut App) {}
+
+fn app() -> App {
+    let mut app = internal_test_util::create_minimal_app();
+    app.add_plugins(StatesPlugin);
+    app.add_plugins(plugin);
+    app
+}
+
+#[test]
+fn test_auto_init_state() {
+    let app = app();
+    assert!(
+        app.world().get_resource::<State<Test>>().is_some(),
+        "did not auto init state"
+    );
+    assert!(
+        app.world().get_resource::<NextState<Test>>().is_some(),
+        "did not auto init state"
+    );
+    assert!(
+        app.world().get_resource::<State<InnerTest>>().is_some(),
+        "did not auto init state"
+    );
+    assert!(
+        app.world().get_resource::<NextState<InnerTest>>().is_some(),
+        "did not auto init state"
+    );
+}