statum-macros 0.8.9

Proc macros for representing legal workflow and protocol states explicitly in Rust
Documentation
use quote::{format_ident, quote};

use super::{EnumInfo, ParsedVariantShape};

pub fn generate_state_impls(enum_info: &EnumInfo) -> proc_macro2::TokenStream {
    let state_trait_ident = enum_info.get_trait_name();
    let parsed_enum = match enum_info.parse() {
        Ok(parsed) => parsed,
        Err(err) => return err,
    };
    let vis = parsed_enum.vis;
    let derive_tokens = parsed_enum
        .derives
        .iter()
        .map(quote::ToTokens::to_token_stream)
        .collect::<Vec<_>>();

    let mut variant_structs = Vec::with_capacity(enum_info.variants.len());
    for variant in parsed_enum.variants {
        let variant_name = format_ident!("{}", variant.name);
        let variant_derives = if derive_tokens.is_empty() {
            quote! {}
        } else {
            quote! { #[derive(#(#derive_tokens),*)] }
        };

        let tokens = match &variant.shape {
            ParsedVariantShape::Unit => {
                quote! {
                    #variant_derives
                    #vis struct #variant_name;

                    impl #state_trait_ident for #variant_name {
                        type Data = ();
                    }

                    impl statum::StateMarker for #variant_name {
                        type Data = ();
                    }

                    impl statum::UnitState for #variant_name {}
                }
            }
            ParsedVariantShape::Tuple { data_type } => {
                quote! {
                    #variant_derives
                    #vis struct #variant_name (pub #data_type);

                    impl #state_trait_ident for #variant_name {
                        type Data = #data_type;
                    }

                    impl statum::StateMarker for #variant_name {
                        type Data = #data_type;
                    }

                    impl statum::DataState for #variant_name {}
                }
            }
            ParsedVariantShape::Named {
                data_struct_ident,
                fields,
            } => {
                let payload_fields = fields.iter().map(|field| {
                    let field_ident = &field.ident;
                    let field_type = &field.field_type;
                    quote! { pub #field_ident: #field_type }
                });

                quote! {
                    #variant_derives
                    #vis struct #data_struct_ident {
                        #(#payload_fields),*
                    }

                    #variant_derives
                    #vis struct #variant_name (pub #data_struct_ident);

                    impl #state_trait_ident for #variant_name {
                        type Data = #data_struct_ident;
                    }

                    impl statum::StateMarker for #variant_name {
                        type Data = #data_struct_ident;
                    }

                    impl statum::DataState for #variant_name {}
                }
            }
        };
        variant_structs.push(tokens);
    }

    let state_trait = quote! {
        #enum_info
        #vis trait #state_trait_ident {
            type Data;
        }
    };

    let uninitialized_state_name = format_ident!("Uninitialized{}", enum_info.name);

    let uninitialized_state = quote! {
        pub struct #uninitialized_state_name;

        impl #state_trait_ident for #uninitialized_state_name {
            type Data = ();
        }

        impl statum::StateMarker for #uninitialized_state_name {
            type Data = ();
        }

        impl statum::UnitState for #uninitialized_state_name {}
    };

    quote! {
        #state_trait

        #(#variant_structs)*

        #uninitialized_state
    }
}