buf_sized_derive/
lib.rs

1//! A library for deriving the `BufSized` trait.
2//!
3//! # Examples
4//! ```
5//! use buf_sized_derive::BufSized;
6//!
7//! #[derive(BufSized)]
8//! struct Data {
9//!     header: u8,
10//!     num: u128,
11//!     payload: [u8; 12],
12//!     crc: u32,
13//! }
14//! ```
15use 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/// Derive the `BufSized` trait for a struct.
23#[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}