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}