use proc_macro2::TokenStream;

use quote::{quote, quote_spanned};

use syn::spanned::Spanned;

use syn::{

Data, DeriveInput, Fields, GenericParam, Generics, Index, parse_macro_input, parse_quote,

};

#[proc_macro_derive(HeapSize)]

pub fn derive_heap_size(input: proc_macro::TokenStream) -> proc_macro::TokenStream {

// Parse the input tokens into a syntax tree.

let input = parse_macro_input!(input as DeriveInput);

// Used in the quasi-quotation below as `#name`.

let name = input.ident;

// Add a bound `T: HeapSize` to every type parameter T.

let mut generics = input.generics;

add_trait_bounds(&mut generics);

let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

// Generate an expression to sum up the heap size of each field.

let sum = heap_size_sum(&input.data);

// The generated impl.

let expanded = quote! {

impl #impl_generics ::heapsize::HeapSize for #name #ty_generics #where_clause {

fn heap_size_of_children(&self) -> usize {

#sum

}

}

};

// Hand the output tokens back to the compiler.

proc_macro::TokenStream::from(expanded)

}

// Add a bound `T: HeapSize` to every type parameter T.

fn add_trait_bounds(generics: &mut Generics) {

for param in &mut generics.params {

if let GenericParam::Type(type_param) = param {

type_param.bounds.push(parse_quote!(::heapsize::HeapSize));

}

}

}

// Generate an expression to sum up the heap size of each field.

fn heap_size_sum(data: &Data) -> TokenStream {

match data {

Data::Struct(data) => {

match &data.fields {

Fields::Named(fields) => {

// Expands to an expression like

//

// 0 + self.x.heap_size() + self.y.heap_size() + self.z.heap_size()

//

// but using fully qualified function call syntax.

//

// We take some care to use the span of each `syn::Field` as

// the span of the corresponding `heap_size_of_children`

// call. This way if one of the field types does not

// implement `HeapSize` then the compiler's error message

// underlines which field it is. An example is shown in the

// readme of the parent directory.

let recurse = fields.named.iter().map(|f| {

let name = &f.ident;

let field_access = quote_spanned!(f.span()=> &self.#name);

quote! {

::heapsize::HeapSize::heap_size_of_children(#field_access)

}

});

quote! {

0 #(+ #recurse)*

}

}

Fields::Unnamed(fields) => {

// Expands to an expression like

//

// 0 + self.0.heap_size() + self.1.heap_size() + self.2.heap_size()

let recurse = fields.unnamed.iter().enumerate().map(|(i, f)| {

let index = Index::from(i);

let field_access = quote_spanned!(f.span()=> &self.#index);

quote! {

::heapsize::HeapSize::heap_size_of_children(#field_access)

}

});

quote! {

0 #(+ #recurse)*

}

}

Fields::Unit => {

// Unit structs cannot own more than 0 bytes of heap memory.

quote!(0)

}

}

}

Data::Enum(_) | Data::Union(_) => unimplemented!(),

}

}