bevy/crates/bevy_ecs/macros/src/states.rs

45 lines
1.3 KiB
Rust
Raw Normal View History

use proc_macro::{Span, TokenStream};
use quote::{format_ident, quote};
use syn::{parse_macro_input, Data::Enum, DeriveInput};
use crate::bevy_ecs_path;
pub fn derive_states(input: TokenStream) -> TokenStream {
let ast = parse_macro_input!(input as DeriveInput);
let error = || {
syn::Error::new(
Span::call_site().into(),
"derive(States) only supports fieldless enums",
)
.into_compile_error()
.into()
};
let Enum(enumeration) = ast.data else {
return error();
};
if enumeration.variants.iter().any(|v| !v.fields.is_empty()) {
return error();
}
let generics = ast.generics;
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
let mut trait_path = bevy_ecs_path();
trait_path.segments.push(format_ident!("schedule").into());
trait_path.segments.push(format_ident!("States").into());
let struct_name = &ast.ident;
let idents = enumeration.variants.iter().map(|v| &v.ident);
let len = idents.len();
quote! {
impl #impl_generics #trait_path for #struct_name #ty_generics #where_clause {
type Iter = std::array::IntoIter<Self, #len>;
fn variants() -> Self::Iter {
[#(Self::#idents,)*].into_iter()
}
}
}
.into()
}