1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
extern crate proc_macro;

use oasis_borsh_derive_internal::*;
use proc_macro::TokenStream;
use quote::{quote, format_ident};
use syn::{parse_macro_input, Ident, ItemEnum, ItemStruct, ItemUnion, LitInt, Token};

#[proc_macro_derive(BorshSerialize, attributes(borsh_skip))]
pub fn borsh_serialize(input: TokenStream) -> TokenStream {
    let res = if let Ok(input) = syn::parse::<ItemStruct>(input.clone()) {
        struct_ser(&input)
    } else if let Ok(input) = syn::parse::<ItemEnum>(input.clone()) {
        enum_ser(&input)
    } else if let Ok(input) = syn::parse::<ItemUnion>(input.clone()) {
        union_ser(&input)
    } else {
        // Derive macros can only be defined on structs, enums, and unions.
        unreachable!()
    };
    TokenStream::from(match res {
        Ok(res) => res,
        Err(err) => err.to_compile_error(),
    })
}

#[proc_macro_derive(BorshDeserialize, attributes(borsh_skip, borsh_init))]
pub fn borsh_deserialize(input: TokenStream) -> TokenStream {
    let res = if let Ok(input) = syn::parse::<ItemStruct>(input.clone()) {
        struct_de(&input)
    } else if let Ok(input) = syn::parse::<ItemEnum>(input.clone()) {
        enum_de(&input)
    } else if let Ok(input) = syn::parse::<ItemUnion>(input.clone()) {
        union_de(&input)
    } else {
        // Derive macros can only be defined on structs, enums, and unions.
        unreachable!()
    };
    TokenStream::from(match res {
        Ok(res) => res,
        Err(err) => err.to_compile_error(),
    })
}

struct SeqMacroSpec {
    mac_ident: Ident,
    prefix: Option<Ident>,
    lengths: Vec<usize>,
}

impl syn::parse::Parse for SeqMacroSpec {
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
        let mac_ident = input.parse()?;
        input.parse::<Token![=>]>()?;

        let prefix = input.parse::<Option<Ident>>()?;
        if prefix.is_some() {
            input.parse::<Token![::]>()?;
        }

        let seq_lengths;
        syn::parenthesized!(seq_lengths in input);
        let punctuated_lengths: syn::punctuated::Punctuated<LitInt, syn::parse::Nothing> =
            seq_lengths.parse_terminated(LitInt::parse)?;
        let lengths = punctuated_lengths
            .iter()
            .map(|l| l.base10_parse::<usize>())
            .collect::<Result<Vec<_>, _>>()?;

        Ok(SeqMacroSpec {
            mac_ident,
            prefix,
            lengths,
        })
    }
}

#[proc_macro]
#[doc(hidden)]
pub fn _gen_seq_macro(input: TokenStream) -> TokenStream {
    let SeqMacroSpec {
        mac_ident,
        prefix,
        lengths,
    } = parse_macro_input!(input as SeqMacroSpec);
    let seqs = lengths.iter().copied().map(|len| {
        (0..len)
            .map(|i| {
                if let Some(prefix) = &prefix {
                    let seq_ident = format_ident!("{}{}", prefix, i);
                    quote!(#seq_ident)
                } else {
                    let seq_lit = proc_macro2::Literal::usize_unsuffixed(i);
                    quote!(#seq_lit)
                }
            })
            .collect::<Vec<_>>()
    });
    let length_lits = lengths
        .iter()
        .map(|&i| proc_macro2::Literal::usize_unsuffixed(i));
    let ts = TokenStream::from(quote! {
        #mac_ident!(#(#length_lits => ( #(#seqs)* ))*);
    });
    ts
}