bitcoin-consensus-derive 0.2.2

///! proc-macro to derive a bitcoin `Encodable` and `Decodable` implementation for a struct
use proc_macro::TokenStream;
use proc_macro2::{Ident, TokenStream as TokenStream2};
use quote::{quote, ToTokens};
use syn::{
    parse_macro_input, Data, DataStruct, DeriveInput, Fields, FieldsNamed, FieldsUnnamed, Index,
    Type,
};

/// Derive `Encodable` for a struct.
///
/// Notes:
/// - all number fields will be encoded in big endian, unlike rust-bitcoin
/// - all `Option` fields will be encoded with a `bool` indicating whether the field is `Some` or `None`
#[proc_macro_derive(Encodable)]
pub fn derive_encodable(input: TokenStream) -> TokenStream {
    let DeriveInput { ident, data, .. } = parse_macro_input!(input);
    // handle struct
    let output = if let Data::Struct(DataStruct {
        fields: Fields::Named(FieldsNamed { named: fields, .. }),
        ..
    }) = data
    {
        let field_tokens = fields.iter().map(|field| {
            let field_name = field.ident.as_ref().unwrap();
            let field_type = &field.ty;
            generate_field_encode(true, field_name, field_type)
        });
        let output = quote! {
            impl bitcoin::consensus::Encodable for #ident {
                fn consensus_encode<W: bitcoin::io::Write + ?Sized>(
                    &self,
                    w: &mut W,
                ) -> core::result::Result<usize, bitcoin::io::Error> {
                    let mut len = 0;
                    #( #field_tokens )*
                    Ok(len)
                }
            }
        };
        output
    } else if let Data::Struct(DataStruct {
        fields: Fields::Unnamed(FieldsUnnamed {
            unnamed: fields, ..
        }),
        ..
    }) = data
    {
        let field_tokens = fields.iter().enumerate().map(|(i, field)| {
            let field_name = Index::from(i);
            let field_type = &field.ty;
            generate_field_encode(true, &field_name, field_type)
        });
        let output = quote! {
            impl bitcoin::consensus::Encodable for #ident {
                fn consensus_encode<W: bitcoin::io::Write + ?Sized>(
                    &self,
                    w: &mut W,
                ) -> core::result::Result<usize, bitcoin::io::Error> {
                    let mut len = 0;
                    #( #field_tokens )*
                    Ok(len)
                }
            }
        };
        output
    } else {
        unimplemented!()
    };
    output.into()
}

/// Derive `Decodable` for a struct.
///
/// See [`derive_encodable`] for notes.
#[proc_macro_derive(Decodable)]
pub fn derive_decodable(input: TokenStream) -> TokenStream {
    let DeriveInput { ident, data, .. } = parse_macro_input!(input);
    // handle struct
    if let syn::Data::Struct(syn::DataStruct {
        fields: syn::Fields::Named(FieldsNamed { named: fields, .. }),
        ..
    }) = data
    {
        let field_tokens = fields.iter().map(|field| {
            let field_name = field.ident.as_ref().unwrap();
            let field_type = &field.ty;
            generate_field_decode(field_name, field_type)
        });
        let field_names = fields.iter().map(|field| {
            let field_name = field.ident.as_ref().unwrap();
            quote! {
                #field_name,
            }
        });
        let output = quote! {
            impl bitcoin::consensus::Decodable for #ident {
                fn consensus_decode<R: bitcoin::io::Read + ?Sized>(
                    r: &mut R,
                ) -> core::result::Result<Self, bitcoin::consensus::encode::Error> {
                    #( #field_tokens )*
                    Ok(Self {
                        #( #field_names )*
                    })
                }
            }
        };
        return output.into();
    } else if let Data::Struct(DataStruct {
        fields: Fields::Unnamed(FieldsUnnamed {
            unnamed: fields, ..
        }),
        ..
    }) = data
    {
        let field_tokens = fields.iter().enumerate().map(|(i, field)| {
            let field_name = Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
            let field_type = &field.ty;
            generate_field_decode(&field_name, field_type)
        });
        let field_names = fields.iter().enumerate().map(|(i, _field)| {
            let field_name = Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
            quote! {
                #field_name,
            }
        });
        let output = quote! {
            impl bitcoin::consensus::Decodable for #ident {
                fn consensus_decode<R: bitcoin::io::Read + ?Sized>(
                    r: &mut R,
                ) -> core::result::Result<Self, bitcoin::consensus::encode::Error> {
                    #( #field_tokens )*
                    Ok(Self(
                        #( #field_names )*
                    ))
                }
            }
        };
        return output.into();
    } else {
        unimplemented!()
    }
}

fn generate_field_encode(
    is_self: bool,
    field_name: &dyn ToTokens,
    field_type: &Type,
) -> TokenStream2 {
    let field_access = if is_self {
        quote! {
            self.#field_name
        }
    } else {
        quote! {
            #field_name
        }
    };
    if get_array_length(field_type).is_some() {
        quote! {
            for el in &#field_access {
                len += el.consensus_encode(w)?;
            }
        }
    } else if is_numeric_type(field_type) {
        quote! {
            let buf = #field_access.to_be_bytes();
            w.write_all(&buf)?;
            len += buf.len();
        }
    } else if let Some(inner_type) = extract_option_type(field_type) {
        let inner_tokens = generate_field_encode(
            false,
            &Ident::new("inner", proc_macro2::Span::call_site()),
            inner_type,
        );
        quote! {
            len += #field_access.is_some().consensus_encode(w)?;
            if let Some(inner) = &#field_access {
                #inner_tokens
            }
        }
    } else {
        quote! {
            len += #field_access.consensus_encode(w)?;
        }
    }
}

fn generate_field_decode(var: &Ident, field_type: &Type) -> TokenStream2 {
    let output = if let Some(size) = get_array_length(field_type) {
        quote! {
            use core::convert::TryInto;
            use alloc::vec::Vec;
            let mut v = Vec::with_capacity(#size);
            for _ in 0..#size {
                let el = bitcoin::consensus::Decodable::consensus_decode(r)?;
                v.push(el);
            }
            let #var = v.try_into().unwrap();
        }
    } else if is_numeric_type(field_type) {
        quote! {
            let mut buf = [0u8; core::mem::size_of::<#field_type>()];
            r.read_exact(&mut buf)?;
            let #var = #field_type::from_be_bytes(buf);
        }
    } else if let Some(inner_type) = extract_option_type(field_type) {
        let inner_tokens = generate_field_decode(
            &Ident::new("inner", proc_macro2::Span::call_site()),
            inner_type,
        );
        quote! {
            let is_some: bool = bitcoin::consensus::Decodable::consensus_decode(r)?;
            let #var = if is_some {
                let inner = {
                    #inner_tokens
                    inner
                };
                Some(inner)
            } else {
                None
            };
        }
    } else {
        quote! {
            let #var = bitcoin::consensus::Decodable::consensus_decode(r)?;
        }
    };
    output
}

fn extract_option_type(ty: &syn::Type) -> Option<&syn::Type> {
    if let syn::Type::Path(syn::TypePath {
        path: syn::Path { segments, .. },
        ..
    }) = ty
    {
        if let Some(syn::PathSegment {
            ident,
            arguments:
                syn::PathArguments::AngleBracketed(syn::AngleBracketedGenericArguments { args, .. }),
            ..
        }) = segments.first()
        {
            if ident == "Option" {
                if let Some(syn::GenericArgument::Type(inner_type)) = args.first() {
                    return Some(inner_type);
                }
            }
        }
    }
    None
}

fn is_numeric_type(ty: &syn::Type) -> bool {
    if let syn::Type::Path(syn::TypePath {
        path: syn::Path { segments, .. },
        ..
    }) = ty
    {
        if let Some(syn::PathSegment { ident, .. }) = segments.first() {
            if ident == "u8"
                || ident == "u16"
                || ident == "u32"
                || ident == "u64"
                || ident == "u128"
                || ident == "i8"
                || ident == "i16"
                || ident == "i32"
                || ident == "i64"
                || ident == "i128"
            {
                return true;
            }
        }
    }
    false
}

fn get_array_length(ty: &syn::Type) -> Option<usize> {
    if let syn::Type::Array(syn::TypeArray { len, .. }) = ty {
        if let syn::Expr::Lit(syn::ExprLit {
            lit: syn::Lit::Int(int),
            ..
        }) = len
        {
            if let Ok(value) = int.base10_parse::<usize>() {
                return Some(value);
            }
        }
    }
    None
}