Skip to main content

borsh_max_size_derive/
lib.rs

1use proc_macro2::TokenStream;
2use quote::{quote, quote_spanned};
3use syn::{
4    parse_macro_input, parse_quote, spanned::Spanned, Data, DeriveInput, Fields, GenericParam,
5    Generics,
6};
7
8extern crate proc_macro;
9
10#[proc_macro_derive(MaxSize)]
11pub fn derive_max_size(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
12    let mut input = parse_macro_input!(input as DeriveInput);
13
14    let name = &input.ident;
15
16    // Add a bound `T: MaxSize` to every type parameter T.
17    add_trait_bounds(&mut input.generics);
18    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
19
20    // Generate an expression to aggregate the max size of each field.
21    let sum = overall_max_size(&input.data);
22
23    let expanded = quote! {
24        // The generated impl.
25        impl #impl_generics borsh_max_size::MaxSize for #name #ty_generics #where_clause {
26            fn max_size() -> 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: MaxSize` to every type parameter T.
37fn add_trait_bounds(generics: &mut Generics) {
38    for param in &mut generics.params {
39        if let GenericParam::Type(ref mut type_param) = *param {
40            type_param
41                .bounds
42                .push(parse_quote!(borsh_max_size::MaxSize));
43        }
44    }
45}
46
47fn overall_max_size(data: &Data) -> TokenStream {
48    match *data {
49        Data::Struct(ref data) => size_of_fields(&data.fields),
50        Data::Enum(ref data) => {
51            let recurse = data
52                .variants
53                .iter()
54                .map(|variant| size_of_fields(&variant.fields));
55            recurse.fold(
56                quote!(0),
57                |max_, expr| quote! { std::cmp::max(#max_, #expr)},
58            )
59        }
60        Data::Union(ref data) => {
61            let recurse = data.fields.named.iter().map(|f| {
62                let type_ = &f.ty;
63                quote_spanned! {f.span()=>
64                    <#type_ as borsh_max_size::MaxSize>::max_size()
65                }
66            });
67            recurse.fold(
68                quote!(0),
69                |max_, expr| quote! { std::cmp::max(#max_, #expr)},
70            )
71        }
72    }
73}
74
75fn size_of_fields(fields: &Fields) -> TokenStream {
76    match fields {
77        Fields::Named(ref fields) => {
78            // Expands to an expression like
79            //
80            //     0 + self.x.max_size() + self.y.max_size() + self.z.max_size()
81            //
82            // but using fully qualified function call syntax.
83            let recurse = fields.named.iter().map(|f| {
84                let type_ = &f.ty;
85                quote_spanned! {f.span()=>
86                    <#type_ as borsh_max_size::MaxSize>::max_size()
87                }
88            });
89            quote! {
90                0 #(+ #recurse)*
91            }
92        }
93        Fields::Unnamed(ref fields) => {
94            // Expands to an expression like
95            //
96            //     0 + self.0.max_size() + self.1.max_size() + self.2.max_size()
97            let recurse = fields.unnamed.iter().map(|f| {
98                let type_ = &f.ty;
99                quote_spanned! {f.span()=>
100                    <#type_ as borsh_max_size::MaxSize>::max_size()
101                }
102            });
103            quote! {
104                0 #(+ #recurse)*
105            }
106        }
107        Fields::Unit => {
108            // Unit structs don't take up space in the serialization.
109            quote!(0)
110        }
111    }
112}