Rand Derive Macro

lib.rs

#[derive(FromVariant)]
struct EnumVariantData {
    ident: syn::Ident,
}

#[derive(FromDeriveInput)]
#[darling(supports(enum_unit), forward_attrs(allow, doc, cfg))]
struct EnumData {
    ident: syn::Ident,
    data: ast::Data<EnumVariantData, ()>,
}

#[proc_macro_derive(RandomSelect)]
pub fn random_select(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let target = match EnumData::from_derive_input(&input) {
        Ok(parsed) => parsed,
        Err(e) => return e.write_errors().into(),
    };

    let enum_name = target.ident;
    let variant_data = target.data.take_enum().unwrap();
    let num_variants = variant_data.len();

    let mut branches = variant_data.into_iter().enumerate().map(|(i, v)| {
        let name = v.ident;
        quote! {
            #i => #enum_name::#name
        }
    }).collect::<Vec<_>>();

    branches.push(quote! {
        _ => unreachable!()
    });

    let tokens = quote! {
        impl ::rand::distributions::Distribution<#enum_name> for ::rand::distributions::Standard {
            fn sample<R: ::rand::Rng + ?Sized>(&self, rng: &mut R) -> #enum_name {
                match rng.gen_range(0..#num_variants) {
                    #(#branches),*
                }
            }
        }
    };

    tokens.into()
}