collum_derive/
lib.rs

1use darling::{FromDeriveInput, FromMeta};
2use proc_macro2::{Span, TokenStream};
3use quote::{quote, quote_spanned};
4use syn::spanned::Spanned;
5use syn::{parse_macro_input, parse_quote, Field, FieldsNamed, FieldsUnnamed, Type};
6use syn::{Data, DeriveInput, Fields, GenericParam, Generics, Ident};
7
8#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, FromMeta)]
9enum Endian {
10    #[default]
11    Native,
12    Little,
13    Big,
14}
15
16impl Endian {
17    fn transform_captured_type(
18        &self,
19        ident: Option<&Ident>,
20        ty: &Type,
21    ) -> (TokenStream, TokenStream) {
22        match self {
23            Endian::Native => (quote!(#ident), quote!(#ty)),
24            Endian::Little => (
25                quote!(collum::endian::Little(#ident)),
26                quote!(collum::endian::Little<#ty>),
27            ),
28            Endian::Big => (
29                quote!(collum::endian::Big(#ident)),
30                quote!(collum::endian::Big<#ty>),
31            ),
32        }
33    }
34
35    fn field_to_captured_type(&self, field: &Field) -> (TokenStream, TokenStream) {
36        let ident = &field.ident;
37        let ty = &field.ty;
38
39        self.transform_captured_type(ident.as_ref(), ty)
40    }
41}
42
43#[derive(FromDeriveInput, Default)]
44#[darling(default, attributes(collum))]
45struct Opts {
46    #[darling(default)]
47    endian: Endian,
48}
49
50fn emit_impl_body_struct_named(
51    ident: &Ident,
52    accvar: &Ident,
53    fields: &FieldsNamed,
54    opts: &Opts,
55) -> TokenStream {
56    let binding: Vec<_> = fields.named.iter().map(|f| &f.ident).collect();
57
58    let (capturing, ty): (Vec<_>, Vec<_>) = fields
59        .named
60        .iter()
61        .map(|f| opts.endian.field_to_captured_type(f))
62        .unzip();
63
64    quote! {
65        #( let #binding; )*
66
67        #( (#capturing, #accvar) = <#ty as collum::Collum>::take_from_bit_slice(#accvar)?; )*
68
69        Some(( #ident { #( #binding ),* }, #accvar))
70    }
71}
72
73fn emit_impl_body_struct_unnamed(
74    ident: &Ident,
75    accvar: &Ident,
76    fields: &FieldsUnnamed,
77    opts: &Opts,
78) -> TokenStream {
79    let unnamed_to_numbered = |i: usize, f: &Field| -> Ident {
80        let name = format!("v{i}");
81        Ident::new(&name, f.span())
82    };
83
84    let binding: Vec<_> = fields
85        .unnamed
86        .iter()
87        .enumerate()
88        .map(|(i, f)| unnamed_to_numbered(i, f))
89        .collect();
90
91    let (capturing, ty): (Vec<_>, Vec<_>) = fields
92        .unnamed
93        .iter()
94        .enumerate()
95        .map(|(i, f)| {
96            let named = unnamed_to_numbered(i, f);
97            opts.endian.transform_captured_type(Some(&named), &f.ty)
98        })
99        .unzip();
100
101    quote! {
102        #( let #binding; )*
103
104        #( (#capturing, #accvar) = <#ty as collum::Collum>::take_from_bit_slice(#accvar)?; )*
105
106        Some((#ident( #( #binding ),* ), #accvar))
107    }
108}
109
110fn emit_impl_body(ident: &Ident, accvar: &Ident, data: &Data, opts: &Opts) -> TokenStream {
111    match *data {
112        Data::Struct(ref data) => match &data.fields {
113            Fields::Named(fields) => emit_impl_body_struct_named(ident, accvar, fields, opts),
114            Fields::Unnamed(fields) => emit_impl_body_struct_unnamed(ident, accvar, fields, opts),
115            Fields::Unit => {
116                quote!(Some((#ident, #accvar)))
117            }
118        },
119        Data::Enum(..) | Data::Union(..) => unimplemented!(),
120    }
121}
122
123fn impl_collum(ast: DeriveInput, opts: Opts) -> TokenStream {
124    let name = ast.ident;
125    let slice_var = Ident::new("slice", Span::call_site());
126    let generics = add_collum_sized_bound(ast.generics);
127
128    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
129
130    let impl_body = emit_impl_body(&name, &slice_var, &ast.data, &opts);
131
132    let expanded = quote! {
133        impl #impl_generics Collum for #name #ty_generics #where_clause {
134            fn take_from_bit_slice(mut slice: collum::BitSlice) -> Option<(Self, collum::BitSlice)> {
135                #impl_body
136            }
137        }
138    };
139
140    expanded
141}
142
143fn impl_collum_sized(ast: DeriveInput) -> TokenStream {
144    let name = ast.ident;
145    let generics = add_collum_sized_bound(ast.generics);
146
147    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
148
149    fn compute_size(data: &Data) -> TokenStream {
150        match *data {
151            Data::Struct(ref data) => {
152                let size = data.fields.iter().map(|f| {
153                    let ty = &f.ty;
154                    quote_spanned! { f.span() =>
155                        <#ty as collum::Sized>::bits()
156                    }
157                });
158
159                quote! {
160                    0 #(+ #size)*
161                }
162            }
163            Data::Enum(..) | Data::Union(..) => unimplemented!(),
164        }
165    }
166
167    let size = compute_size(&ast.data);
168
169    let expanded = quote! {
170        impl #impl_generics collum::Sized for #name #ty_generics #where_clause {
171            fn bits() -> usize {
172                #size
173            }
174        }
175    };
176
177    expanded
178}
179
180fn add_collum_sized_bound(mut generics: Generics) -> Generics {
181    for param in &mut generics.params {
182        if let GenericParam::Type(ref mut type_param) = *param {
183            type_param.bounds.push(parse_quote!(collum::Sized))
184        }
185    }
186    generics
187}
188
189#[proc_macro_derive(Collum, attributes(collum))]
190pub fn collum_derive(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
191    let ast = parse_macro_input!(item as DeriveInput);
192    let opts = Opts::from_derive_input(&ast).expect("Wrong options");
193
194    proc_macro::TokenStream::from(impl_collum(ast, opts))
195}
196
197#[proc_macro_derive(Sized)]
198pub fn collum_sized_derive(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
199    let ast = parse_macro_input!(item as DeriveInput);
200
201    proc_macro::TokenStream::from(impl_collum_sized(ast))
202}