deepsize_derive/
lib.rs

1//! A basic [`DeepSizeOf`](deepsize::DeepSizeOf) Derive implementation
2//!
3//! Mainly from `syn`'s [`heap_size` derive example][heap_size]
4//!
5//! [heap_size]: https://github.com/dtolnay/syn/commits/master/examples/heapsize/heapsize_derive/src/lib.rs
6
7extern crate proc_macro;
8
9use proc_macro2::TokenStream;
10use quote::{quote, quote_spanned};
11use syn::spanned::Spanned;
12use syn::{
13    parse_macro_input, parse_quote, Data, DeriveInput, Fields, GenericParam, Generics, Index,
14};
15
16#[proc_macro_derive(DeepSizeOf)]
17pub fn derive_deep_size(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
18    // Parse the input tokens into a syntax tree.
19    let input = parse_macro_input!(input as DeriveInput);
20
21    // Used in the quasi-quotation below as `#name`.
22    let name = input.ident;
23
24    // Add a bound `T: DeepSizeOf` to every type parameter T.
25    let generics = add_trait_bounds(input.generics);
26    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
27
28    // Generate an expression to sum up the size of each field.
29    let sum = deepsize_sum(&input.data, &name);
30
31    let expanded = quote! {
32        // The generated impl.
33        impl #impl_generics ::deepsize::DeepSizeOf for #name #ty_generics #where_clause {
34            fn deep_size_of_children(&self, context: &mut ::deepsize::Context) -> usize {
35                #sum
36            }
37        }
38    };
39
40    // Hand the output tokens back to the compiler.
41    proc_macro::TokenStream::from(expanded)
42}
43
44// Add a bound `T: DeepSizeOf` to every type parameter T.
45fn add_trait_bounds(mut generics: Generics) -> Generics {
46    for param in &mut generics.params {
47        if let GenericParam::Type(ref mut type_param) = *param {
48            type_param.bounds.push(parse_quote!(::deepsize::DeepSizeOf));
49        }
50    }
51    generics
52}
53
54fn match_fields(fields: &syn::Fields) -> TokenStream {
55    match fields {
56        Fields::Named(ref fields) => {
57            let recurse = fields.named.iter().map(|f| {
58                let name = &f.ident;
59                quote_spanned! {f.span()=>
60                    ::deepsize::DeepSizeOf::deep_size_of_children(&self.#name, context)
61                }
62            });
63            quote! {
64                0 #(+ #recurse)*
65            }
66        }
67        Fields::Unnamed(ref fields) => {
68            let recurse = fields.unnamed.iter().enumerate().map(|(i, f)| {
69                let index = Index::from(i);
70                quote_spanned! {f.span()=>
71                    ::deepsize::DeepSizeOf::deep_size_of_children(&self.#index, context)
72                }
73            });
74            quote! {
75                0 #(+ #recurse)*
76            }
77        }
78        Fields::Unit => {
79            // Unit structs cannot own more than 0 bytes of memory.
80            quote!(0)
81        }
82    }
83}
84
85fn match_enum_fields(fields: &syn::Fields) -> TokenStream {
86    match fields {
87        Fields::Named(ref fields) => {
88            let recurse = fields.named.iter().map(|f| {
89                let name = &f.ident;
90                quote_spanned! {f.span()=>
91                    ::deepsize::DeepSizeOf::deep_size_of_children(#name, context)
92                }
93            });
94            quote! {
95                0 #(+ #recurse)*
96            }
97        }
98        Fields::Unnamed(ref fields) => {
99            let recurse = fields.unnamed.iter().enumerate().map(|(i, f)| {
100                let i = syn::Ident::new(&format!("_{}", i), proc_macro2::Span::call_site());
101                quote_spanned! {f.span()=>
102                    ::deepsize::DeepSizeOf::deep_size_of_children(#i, context)
103                }
104            });
105            quote! {
106                0 #(+ #recurse)*
107            }
108        }
109        Fields::Unit => {
110            // Unit structs cannot own more than 0 bytes of memory.
111            quote!(0)
112        }
113    }
114}
115
116fn get_matcher(var: &syn::Variant) -> TokenStream {
117    let matcher = match &var.fields {
118        Fields::Unit => TokenStream::new(),
119        Fields::Unnamed(fields) => {
120            let fields: TokenStream = (0..fields.unnamed.len())
121                .map(|n| {
122                    let i = syn::Ident::new(&format!("_{}", n), proc_macro2::Span::call_site());
123                    quote!(#i,)
124                })
125                .collect();
126            quote!((#fields))
127        }
128        Fields::Named(fields) => {
129            let fields: TokenStream = fields
130                .named
131                .iter()
132                .map(|f| {
133                    let i = f.ident.as_ref().unwrap();
134                    quote!(#i,)
135                })
136                .collect();
137            quote!({#fields})
138        }
139    };
140
141    quote!(#matcher)
142}
143
144/// Generate an expression to sum up the size of each field.
145fn deepsize_sum(data: &Data, struct_name: &proc_macro2::Ident) -> TokenStream {
146    match *data {
147        Data::Struct(ref inner) => match_fields(&inner.fields),
148        Data::Enum(ref inner) => {
149            let arms = inner.variants.iter().map(|var| {
150                let matcher = get_matcher(var);
151                let output = match_enum_fields(&var.fields);
152                let name = &var.ident;
153                let ident = quote!(#struct_name::#name);
154                quote!(#ident #matcher => #output,)
155            });
156
157            quote! {
158                match self {
159                    #(#arms)*
160                    _ => 0 // This is needed for empty enums
161                }
162            }
163        }
164        Data::Union(_) => unimplemented!(),
165    }
166}