bitcoin_consensus_derive/
lib.rs

1///! proc-macro to derive a bitcoin `Encodable` and `Decodable` implementation for a struct
2use proc_macro::TokenStream;
3use proc_macro2::{Ident, TokenStream as TokenStream2};
4use quote::{quote, ToTokens};
5use syn::{
6    parse_macro_input, Data, DataStruct, DeriveInput, Fields, FieldsNamed, FieldsUnnamed, Index,
7    Type,
8};
9
10/// Derive `Encodable` for a struct.
11///
12/// Notes:
13/// - all number fields will be encoded in big endian, unlike rust-bitcoin
14/// - all `Option` fields will be encoded with a `bool` indicating whether the field is `Some` or `None`
15#[proc_macro_derive(Encodable)]
16pub fn derive_encodable(input: TokenStream) -> TokenStream {
17    let DeriveInput { ident, data, .. } = parse_macro_input!(input);
18    // handle struct
19    let output = if let Data::Struct(DataStruct {
20        fields: Fields::Named(FieldsNamed { named: fields, .. }),
21        ..
22    }) = data
23    {
24        let field_tokens = fields.iter().map(|field| {
25            let field_name = field.ident.as_ref().unwrap();
26            let field_type = &field.ty;
27            generate_field_encode(true, field_name, field_type)
28        });
29        let output = quote! {
30            impl bitcoin::consensus::Encodable for #ident {
31                fn consensus_encode<W: bitcoin::io::Write + ?Sized>(
32                    &self,
33                    w: &mut W,
34                ) -> core::result::Result<usize, bitcoin::io::Error> {
35                    let mut len = 0;
36                    #( #field_tokens )*
37                    Ok(len)
38                }
39            }
40        };
41        output
42    } else if let Data::Struct(DataStruct {
43        fields: Fields::Unnamed(FieldsUnnamed {
44            unnamed: fields, ..
45        }),
46        ..
47    }) = data
48    {
49        let field_tokens = fields.iter().enumerate().map(|(i, field)| {
50            let field_name = Index::from(i);
51            let field_type = &field.ty;
52            generate_field_encode(true, &field_name, field_type)
53        });
54        let output = quote! {
55            impl bitcoin::consensus::Encodable for #ident {
56                fn consensus_encode<W: bitcoin::io::Write + ?Sized>(
57                    &self,
58                    w: &mut W,
59                ) -> core::result::Result<usize, bitcoin::io::Error> {
60                    let mut len = 0;
61                    #( #field_tokens )*
62                    Ok(len)
63                }
64            }
65        };
66        output
67    } else {
68        unimplemented!()
69    };
70    output.into()
71}
72
73/// Derive `Decodable` for a struct.
74///
75/// See [`derive_encodable`] for notes.
76#[proc_macro_derive(Decodable)]
77pub fn derive_decodable(input: TokenStream) -> TokenStream {
78    let DeriveInput { ident, data, .. } = parse_macro_input!(input);
79    // handle struct
80    if let syn::Data::Struct(syn::DataStruct {
81        fields: syn::Fields::Named(FieldsNamed { named: fields, .. }),
82        ..
83    }) = data
84    {
85        let field_tokens = fields.iter().map(|field| {
86            let field_name = field.ident.as_ref().unwrap();
87            let field_type = &field.ty;
88            generate_field_decode(field_name, field_type)
89        });
90        let field_names = fields.iter().map(|field| {
91            let field_name = field.ident.as_ref().unwrap();
92            quote! {
93                #field_name,
94            }
95        });
96        let output = quote! {
97            impl bitcoin::consensus::Decodable for #ident {
98                fn consensus_decode<R: bitcoin::io::Read + ?Sized>(
99                    r: &mut R,
100                ) -> core::result::Result<Self, bitcoin::consensus::encode::Error> {
101                    #( #field_tokens )*
102                    Ok(Self {
103                        #( #field_names )*
104                    })
105                }
106            }
107        };
108        return output.into();
109    } else if let Data::Struct(DataStruct {
110        fields: Fields::Unnamed(FieldsUnnamed {
111            unnamed: fields, ..
112        }),
113        ..
114    }) = data
115    {
116        let field_tokens = fields.iter().enumerate().map(|(i, field)| {
117            let field_name = Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
118            let field_type = &field.ty;
119            generate_field_decode(&field_name, field_type)
120        });
121        let field_names = fields.iter().enumerate().map(|(i, _field)| {
122            let field_name = Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
123            quote! {
124                #field_name,
125            }
126        });
127        let output = quote! {
128            impl bitcoin::consensus::Decodable for #ident {
129                fn consensus_decode<R: bitcoin::io::Read + ?Sized>(
130                    r: &mut R,
131                ) -> core::result::Result<Self, bitcoin::consensus::encode::Error> {
132                    #( #field_tokens )*
133                    Ok(Self(
134                        #( #field_names )*
135                    ))
136                }
137            }
138        };
139        return output.into();
140    } else {
141        unimplemented!()
142    }
143}
144
145fn generate_field_encode(
146    is_self: bool,
147    field_name: &dyn ToTokens,
148    field_type: &Type,
149) -> TokenStream2 {
150    let field_access = if is_self {
151        quote! {
152            self.#field_name
153        }
154    } else {
155        quote! {
156            #field_name
157        }
158    };
159    if get_array_length(field_type).is_some() {
160        quote! {
161            for el in &#field_access {
162                len += el.consensus_encode(w)?;
163            }
164        }
165    } else if is_numeric_type(field_type) {
166        quote! {
167            let buf = #field_access.to_be_bytes();
168            w.write_all(&buf)?;
169            len += buf.len();
170        }
171    } else if let Some(inner_type) = extract_option_type(field_type) {
172        let inner_tokens = generate_field_encode(
173            false,
174            &Ident::new("inner", proc_macro2::Span::call_site()),
175            inner_type,
176        );
177        quote! {
178            len += #field_access.is_some().consensus_encode(w)?;
179            if let Some(inner) = &#field_access {
180                #inner_tokens
181            }
182        }
183    } else {
184        quote! {
185            len += #field_access.consensus_encode(w)?;
186        }
187    }
188}
189
190fn generate_field_decode(var: &Ident, field_type: &Type) -> TokenStream2 {
191    let output = if let Some(size) = get_array_length(field_type) {
192        quote! {
193            use core::convert::TryInto;
194            use alloc::vec::Vec;
195            let mut v = Vec::with_capacity(#size);
196            for _ in 0..#size {
197                let el = bitcoin::consensus::Decodable::consensus_decode(r)?;
198                v.push(el);
199            }
200            let #var = v.try_into().unwrap();
201        }
202    } else if is_numeric_type(field_type) {
203        quote! {
204            let mut buf = [0u8; core::mem::size_of::<#field_type>()];
205            r.read_exact(&mut buf)?;
206            let #var = #field_type::from_be_bytes(buf);
207        }
208    } else if let Some(inner_type) = extract_option_type(field_type) {
209        let inner_tokens = generate_field_decode(
210            &Ident::new("inner", proc_macro2::Span::call_site()),
211            inner_type,
212        );
213        quote! {
214            let is_some: bool = bitcoin::consensus::Decodable::consensus_decode(r)?;
215            let #var = if is_some {
216                let inner = {
217                    #inner_tokens
218                    inner
219                };
220                Some(inner)
221            } else {
222                None
223            };
224        }
225    } else {
226        quote! {
227            let #var = bitcoin::consensus::Decodable::consensus_decode(r)?;
228        }
229    };
230    output
231}
232
233fn extract_option_type(ty: &syn::Type) -> Option<&syn::Type> {
234    if let syn::Type::Path(syn::TypePath {
235        path: syn::Path { segments, .. },
236        ..
237    }) = ty
238    {
239        if let Some(syn::PathSegment {
240            ident,
241            arguments:
242                syn::PathArguments::AngleBracketed(syn::AngleBracketedGenericArguments { args, .. }),
243            ..
244        }) = segments.first()
245        {
246            if ident == "Option" {
247                if let Some(syn::GenericArgument::Type(inner_type)) = args.first() {
248                    return Some(inner_type);
249                }
250            }
251        }
252    }
253    None
254}
255
256fn is_numeric_type(ty: &syn::Type) -> bool {
257    if let syn::Type::Path(syn::TypePath {
258        path: syn::Path { segments, .. },
259        ..
260    }) = ty
261    {
262        if let Some(syn::PathSegment { ident, .. }) = segments.first() {
263            if ident == "u8"
264                || ident == "u16"
265                || ident == "u32"
266                || ident == "u64"
267                || ident == "u128"
268                || ident == "i8"
269                || ident == "i16"
270                || ident == "i32"
271                || ident == "i64"
272                || ident == "i128"
273            {
274                return true;
275            }
276        }
277    }
278    false
279}
280
281fn get_array_length(ty: &syn::Type) -> Option<usize> {
282    if let syn::Type::Array(syn::TypeArray { len, .. }) = ty {
283        if let syn::Expr::Lit(syn::ExprLit {
284            lit: syn::Lit::Int(int),
285            ..
286        }) = len
287        {
288            if let Ok(value) = int.base10_parse::<usize>() {
289                return Some(value);
290            }
291        }
292    }
293    None
294}
bitcoin_consensus_derive/lib.rs

bitcoin_consensus_derive/
lib.rs
