1use proc_macro2::TokenStream;
16use quote::quote;
17use syn::{
18 parse_macro_input, parse_quote, Data, DeriveInput, Field, Fields, GenericParam, Generics, Type,
19 TypeParamBound,
20};
21
22#[proc_macro_derive(BufSized)]
24pub fn buf_sized(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
25 let input = parse_macro_input!(input as DeriveInput);
26 let name = input.ident;
27 let generics = add_trait_bounds(input.generics, &parse_quote!(::buf_sized::BufSized));
28 let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
29 let buf_size = sum_buf_size(input.data);
30 let expanded = quote! {
31 impl #impl_generics ::buf_sized::BufSized for #name #ty_generics #where_clause {
32 const BUF_SIZE: usize = #buf_size;
33 }
34 };
35
36 proc_macro::TokenStream::from(expanded)
37}
38
39fn add_trait_bounds(mut generics: Generics, trait_name: &TypeParamBound) -> Generics {
40 for param in &mut generics.params {
41 if let GenericParam::Type(ref mut type_param) = *param {
42 type_param.bounds.push(trait_name.clone());
43 }
44 }
45
46 generics
47}
48
49fn sum_buf_size(data: Data) -> TokenStream {
50 let mut result = quote! {0usize};
51
52 match data {
53 Data::Struct(structure) => match structure.fields {
54 Fields::Named(fields) => {
55 result.add_fields(fields.named);
56 }
57 Fields::Unit => (),
58 Fields::Unnamed(fields) => {
59 result.add_fields(fields.unnamed);
60 }
61 },
62 Data::Enum(_) | Data::Union(_) => {
63 unimplemented!("Enums and unions are not supported");
64 }
65 }
66
67 result
68}
69
70trait TokenStreamExt {
71 fn add_fields<T>(&mut self, fields: T)
72 where
73 T: IntoIterator<Item = Field>;
74 fn add_field(&mut self, field: Field);
75 fn add_type(&mut self, typ: Type);
76}
77
78impl TokenStreamExt for TokenStream {
79 fn add_fields<T>(&mut self, fields: T)
80 where
81 T: IntoIterator<Item = Field>,
82 {
83 for field in fields {
84 self.add_field(field);
85 }
86 }
87
88 fn add_field(&mut self, field: Field) {
89 self.add_type(field.ty);
90 }
91
92 fn add_type(&mut self, typ: Type) {
93 self.extend(quote! { + <#typ as ::buf_sized::BufSized>::BUF_SIZE });
94 }
95}