size_of_no_padding_derive/
lib.rs

1use proc_macro2::{Ident, Span};
2use quote::quote;
3use syn::{parse_macro_input, parse_quote, Data, DeriveInput, Error, Fields, Index};
4
5/// Calculate the size for every field, add up all the size. call `size_of_no_padding_any(&self)`
6/// will give the size of whole struct without padding in runtime.
7#[proc_macro_derive(SizeOfNoPaddingAny)]
8pub fn derive_size_if_no_padding_any(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
9    let input = parse_macro_input!(input as DeriveInput);
10
11    let struct_name = &input.ident;
12    let generics = &input.generics;
13    let generic_types = &input
14        .generics
15        .params
16        .iter()
17        .filter_map(|p| match p {
18            syn::GenericParam::Type(t) => Some(&t.ident),
19            syn::GenericParam::Const(c) => Some(&c.ident),
20            _ => None,
21        })
22        .collect::<Vec<_>>();
23    let fields = match &input.data {
24        Data::Struct(s) => match &s.fields {
25            Fields::Named(n) => n.named.iter(),
26            Fields::Unnamed(u) => u.unnamed.iter(),
27            Fields::Unit => {
28                return Error::new(
29                    Span::call_site(),
30                    "SizeOfNoPaddingAny not applicable for this type",
31                )
32                .into_compile_error()
33                .into();
34            }
35        },
36        Data::Union(u) => u.fields.named.iter(),
37        Data::Enum(e) => {
38            return Error::new(e.enum_token.span, "SizeOfNoPaddingAny not work for enum")
39                .into_compile_error()
40                .into();
41        }
42    }
43    .enumerate()
44    .map(|(i, f)| {
45        f.ident.as_ref().map(|n| quote! { #n }).unwrap_or({
46            let i = Index::from(i);
47            quote! { #i }
48        })
49    })
50    .collect::<Vec<_>>();
51
52    quote! {
53        impl #generics size_of_no_padding::SizeOfAny for #struct_name <#(#generic_types),*> {
54            fn size_of_no_padding_any(&self) -> usize {
55                let mut size = 0usize;
56
57                #(
58                    size += self.#fields.size_of_no_padding_any();
59                )*
60
61                size
62            }
63        }
64    }
65    .into()
66}
67
68/// Create a shadow struct which as same as original struct except it's marked `#[repr(packed)]`.
69/// call `size_of_no_padding()` method on struct will invoke `std::mem::size_of()` function on shadow one in compile time.
70#[proc_macro_derive(SizeOfNoPadding)]
71pub fn derive_size_of_no_padding(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
72    let mut input = parse_macro_input!(input as DeriveInput);
73    if let Err(e) = remove_attrs(&mut input.data) {
74        return e.into_compile_error().into();
75    }
76    input.attrs.clear();
77    input.attrs.push(parse_quote!(#[derive(Copy, Clone)]));
78    input.attrs.push(parse_quote!(#[repr(C, packed)]));
79    input.attrs.push(parse_quote!(#[allow(dead_code)]));
80
81    let generics = input.generics.clone();
82    let generic_types = &input
83        .generics
84        .params
85        .iter()
86        .filter_map(|p| match p {
87            syn::GenericParam::Type(t) => Some(&t.ident),
88            syn::GenericParam::Const(c) => Some(&c.ident),
89            _ => None,
90        })
91        .collect::<Vec<_>>();
92    let old_ident = input.ident.clone();
93    let new_ident = Ident::new(
94        &(format!("{}SizeOfNoPadding", input.ident)),
95        input.ident.span(),
96    );
97    input.ident = new_ident.clone();
98
99    quote! {
100        #input
101
102        impl #generics size_of_no_padding::SizeOfAny for #old_ident <#(#generic_types),*> {
103            fn size_of_no_padding_any(&self) -> usize {
104                use ::size_of_no_padding::SizeOf;
105                Self::size_of_no_padding()
106            }
107        }
108
109        impl #generics size_of_no_padding::SizeOf for #old_ident <#(#generic_types),*> {
110            fn size_of_no_padding() -> usize {
111                std::mem::size_of::<#new_ident <#(#generic_types),*>>()
112            }
113        }
114    }
115    .into()
116}
117
118fn remove_attrs(data: &mut Data) -> Result<(), Error> {
119    match data {
120        Data::Struct(s) => {
121            match &mut s.fields {
122                Fields::Named(n) => n.named.iter_mut().for_each(|f| f.attrs.clear()),
123                Fields::Unnamed(u) => u.unnamed.iter_mut().for_each(|f| f.attrs.clear()),
124                _ => {}
125            }
126            Ok(())
127        }
128        Data::Union(u) => {
129            u.fields.named.iter_mut().for_each(|f| f.attrs.clear());
130            Ok(())
131        }
132        Data::Enum(e) => Err(Error::new(
133            e.enum_token.span,
134            "SizeOfNoPadding not work for enum",
135        )),
136    }
137}