iter_variants_derive 0.2.0

Derive macros for iterating over all variants of a rust struct/enum
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

use proc_macro as pm;
use proc_macro2::{self as pm2, Span};
use quote::{format_ident, quote};
use syn::{Data, DeriveInput, Fields, PathArguments, Type, parse_macro_input, token::PathSep};

fn do_fields(fields: &Fields, ident: pm2::TokenStream) -> pm2::TokenStream {
    match fields {
        Fields::Unnamed(fields_unnamed) => {
            let vs = (0..fields_unnamed.unnamed.len()).map(|i| format_ident!("v{}", i));
            let mut result = {
                let vs = vs.clone();
                quote! {
                    f(#ident(#(#vs,)*))
                }
            };
            for (field, v) in fields_unnamed.unnamed.iter().zip(vs) {
                let mut ty = field.ty.clone();
                // Option<T> -> Option::<T>
                if let Type::Path(type_path) = &mut ty {
                    if let Some(segment) = type_path.path.segments.last_mut() {
                        if let PathArguments::AngleBracketed(generic_arguments) =
                            &mut segment.arguments
                        {
                            generic_arguments.colon2_token = Some(PathSep {
                                spans: [Span::call_site(); 2],
                            })
                        }
                    }
                }
                result = quote! {
                    <#ty as IterVariants>::iter_variants(|#v| #result)
                };
            }
            result
        }
        Fields::Named(fields_named) => {
            let idents: Vec<_> = fields_named
                .named
                .iter()
                .map(|field| field.ident.clone())
                .collect();
            let mut result = quote! {
                f(#ident {
                    #(#idents,)*
                })
            };
            for field in &fields_named.named {
                let ident = field.ident.clone();
                let mut ty = field.ty.clone();
                // Option<T> -> Option::<T>
                if let Type::Path(type_path) = &mut ty {
                    if let Some(segment) = type_path.path.segments.last_mut() {
                        if let PathArguments::AngleBracketed(generic_arguments) =
                            &mut segment.arguments
                        {
                            generic_arguments.colon2_token = Some(PathSep {
                                spans: [Span::call_site(); 2],
                            })
                        }
                    }
                }
                result = quote! {
                    <#ty as IterVariants>::iter_variants(|#ident| #result)
                }
            }
            result
        }
        Fields::Unit => quote! {
            f(#ident)
        },
    }
}

fn used_types(data: &Data) -> Vec<Type> {
    match data {
        Data::Struct(data_struct) => {
            data_struct.fields.iter()
                .map(|field| field.ty.clone())
                .collect()
        },
        Data::Enum(data_enum) => {
            data_enum.variants.iter()
                .flat_map(|variant| variant.fields.iter())
                .map(|field| field.ty.clone())
                .collect()
        },
        _ => unreachable!(),
    }
}

fn do_count_fields(fields: &Fields) -> pm2::TokenStream {
    let fields_count = fields.iter().map(|field| {
        let ty = &field.ty;
        quote! { <#ty as IterVariants>::iter_variants_count() }
    });
    quote! { (1 #(* #fields_count)*) }
}

#[proc_macro_derive(IterVariants)]
pub fn iter_variants_derive(input: pm::TokenStream) -> pm::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    let ident = input.ident;
    let used_types = used_types(&input.data);

    let output = match &input.data {
        Data::Enum(data_enum) => {
            let arms = data_enum.variants.iter().map(|variant| {
                let variant_ident = &variant.ident;
                do_fields(&variant.fields, quote! {Self::#variant_ident})
            });

            quote! {
                #(#arms;)*
            }
        }
        Data::Struct(data_struct) => do_fields(&data_struct.fields, quote! { #ident }),
        _ => syn::Error::new_spanned(&ident, "`IterVariants` can only be derived for enums or structs")
            .to_compile_error(),
    };
    let variants_count = match &input.data {
        Data::Enum(data_enum) => {
            let counts = data_enum.variants.iter().map(|variant| {
                do_count_fields(&variant.fields)
            });

            quote! { 0 #(+ #counts)* }
        }
        Data::Struct(data_struct) => do_count_fields(&data_struct.fields),
        _ => unreachable!(),
    };

    let (implg, typeg, where_clause) = input.generics.split_for_impl();
    let where_clause = where_clause.map(|w| &w.predicates);
    let expanded = quote! {
        impl #implg IterVariants for #ident #typeg
        where
            #(#used_types: IterVariants<IterVariantsInput = #used_types>,)*
            #where_clause
        {
            type IterVariantsInput = Self;
            #[allow(unused_mut)]
            fn iter_variants<F: FnMut(Self::IterVariantsInput)>(mut f: F) {
                #output
            }

            fn iter_variants_count() -> usize {
                #variants_count
            }
        }
    };

    pm::TokenStream::from(expanded)
}