1 use proc_macro2::TokenStream;
2 use quote::{quote, quote_spanned};
3 use syn::spanned::Spanned;
4 use syn::{
5 parse_macro_input, parse_quote, Data, DeriveInput, Fields, GenericParam, Generics, Index,
6 };
7
8 #[proc_macro_derive(HeapSize)]
derive_heap_size(input: proc_macro::TokenStream) -> proc_macro::TokenStream9 pub fn derive_heap_size(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
10 // Parse the input tokens into a syntax tree.
11 let input = parse_macro_input!(input as DeriveInput);
12
13 // Used in the quasi-quotation below as `#name`.
14 let name = input.ident;
15
16 // Add a bound `T: HeapSize` to every type parameter T.
17 let generics = add_trait_bounds(input.generics);
18 let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
19
20 // Generate an expression to sum up the heap size of each field.
21 let sum = heap_size_sum(&input.data);
22
23 let expanded = quote! {
24 // The generated impl.
25 impl #impl_generics heapsize::HeapSize for #name #ty_generics #where_clause {
26 fn heap_size_of_children(&self) -> usize {
27 #sum
28 }
29 }
30 };
31
32 // Hand the output tokens back to the compiler.
33 proc_macro::TokenStream::from(expanded)
34 }
35
36 // Add a bound `T: HeapSize` to every type parameter T.
add_trait_bounds(mut generics: Generics) -> Generics37 fn add_trait_bounds(mut generics: Generics) -> Generics {
38 for param in &mut generics.params {
39 if let GenericParam::Type(ref mut type_param) = *param {
40 type_param.bounds.push(parse_quote!(heapsize::HeapSize));
41 }
42 }
43 generics
44 }
45
46 // Generate an expression to sum up the heap size of each field.
heap_size_sum(data: &Data) -> TokenStream47 fn heap_size_sum(data: &Data) -> TokenStream {
48 match *data {
49 Data::Struct(ref data) => {
50 match data.fields {
51 Fields::Named(ref fields) => {
52 // Expands to an expression like
53 //
54 // 0 + self.x.heap_size() + self.y.heap_size() + self.z.heap_size()
55 //
56 // but using fully qualified function call syntax.
57 //
58 // We take some care to use the span of each `syn::Field` as
59 // the span of the corresponding `heap_size_of_children`
60 // call. This way if one of the field types does not
61 // implement `HeapSize` then the compiler's error message
62 // underlines which field it is. An example is shown in the
63 // readme of the parent directory.
64 let recurse = fields.named.iter().map(|f| {
65 let name = &f.ident;
66 quote_spanned! {f.span()=>
67 heapsize::HeapSize::heap_size_of_children(&self.#name)
68 }
69 });
70 quote! {
71 0 #(+ #recurse)*
72 }
73 }
74 Fields::Unnamed(ref fields) => {
75 // Expands to an expression like
76 //
77 // 0 + self.0.heap_size() + self.1.heap_size() + self.2.heap_size()
78 let recurse = fields.unnamed.iter().enumerate().map(|(i, f)| {
79 let index = Index::from(i);
80 quote_spanned! {f.span()=>
81 heapsize::HeapSize::heap_size_of_children(&self.#index)
82 }
83 });
84 quote! {
85 0 #(+ #recurse)*
86 }
87 }
88 Fields::Unit => {
89 // Unit structs cannot own more than 0 bytes of heap memory.
90 quote!(0)
91 }
92 }
93 }
94 Data::Enum(_) | Data::Union(_) => unimplemented!(),
95 }
96 }
97