autorand_derive/
lib.rs

1#[macro_use]
2extern crate quote;
3extern crate proc_macro;
4
5use proc_macro::TokenStream;
6use quote::ToTokens;
7use syn::{Data, DataEnum, DataStruct, Field, Fields, Ident, WhereClause};
8
9#[proc_macro_derive(Random)]
10pub fn random(input: TokenStream) -> TokenStream {
11    let input = proc_macro2::TokenStream::from(input);
12
13    let output: proc_macro2::TokenStream = {
14        let parsed = syn::parse2(input).unwrap();
15        impl_random(&parsed)
16    };
17
18    proc_macro::TokenStream::from(output)
19}
20
21fn impl_random(ast: &syn::DeriveInput) -> proc_macro2::TokenStream {
22    let name = &ast.ident;
23    let (impl_generics, ty_generics, ..) = ast.generics.split_for_impl();
24    let type_params: Vec<Ident> = ast
25        .generics
26        .type_params()
27        .map(|t| t.ident.clone())
28        .collect();
29
30    let where_clause = if type_params.is_empty() {
31        ast.generics
32            .where_clause
33            .clone()
34            .map(WhereClause::into_token_stream)
35    } else {
36        let where_clause = ast.generics.where_clause.as_ref();
37        if let Some(where_clause) = where_clause {
38            let predicates = &where_clause.predicates;
39            Some(quote!(
40                where #(#predicates,)* #(#type_params: autorand::Random,)*
41            ))
42        } else {
43            Some(quote!(
44                where #(#type_params: autorand::Random,)*
45            ))
46        }
47    };
48
49    let body = match ast.data {
50        Data::Struct(ref data) => expand_struct_random_body(data),
51        Data::Enum(ref data) => expand_enum_random_body(name, data),
52        Data::Union(_) => panic!("Random derive is not supported for Union types"),
53    };
54
55    let tokens = quote! {
56        impl #impl_generics autorand::Random for #name #ty_generics #where_clause {
57            fn random() -> Self {
58                #body
59            }
60        }
61    };
62
63    //panic!("{}", tokens);
64
65    tokens
66}
67
68fn expand_struct_random_body(data: &DataStruct) -> proc_macro2::TokenStream {
69    let fields = expand_named_fields(data.fields.iter());
70    quote!(
71        Self {
72            #(#fields),*
73        }
74    )
75}
76
77fn expand_named_fields<'a>(
78    fields: impl Iterator<Item = &'a Field> + 'a,
79) -> impl Iterator<Item = impl ToTokens> + 'a {
80    fields.map(|f| {
81        let name = &f.ident;
82        let ty = &f.ty;
83        if f.ident.is_some() {
84            quote! {
85                #name: <#ty as autorand::Random>::random()
86            }
87        } else {
88            quote! {
89                <#ty as autorand::Random>::random()
90            }
91        }
92    })
93}
94
95fn expand_enum_random_body(enum_name: &Ident, data: &DataEnum) -> proc_macro2::TokenStream {
96    let constructors = data.variants.iter().map(|v| {
97        let name = &v.ident;
98        match &v.fields {
99            Fields::Named(fields) => {
100                let fields = expand_named_fields(fields.named.iter());
101                quote!(
102                    #enum_name::#name { #(#fields),* }
103                )
104            }
105            Fields::Unnamed(fields) => {
106                let fields = expand_named_fields(fields.unnamed.iter());
107                quote!(
108                    #enum_name::#name ( #(#fields),* )
109                )
110            }
111            Fields::Unit => quote!(#enum_name::#name),
112        }
113    });
114
115    let matches = constructors
116        .enumerate()
117        .map(|(i, c)| quote!(#i => #c))
118        .collect::<Vec<_>>();
119    let count = matches.len();
120
121    quote!(
122        let variant = autorand::rand::random::<usize>() % #count;
123        match variant {
124            #(#matches),*,
125            _ => unreachable!(),
126        }
127    )
128}