proc_concat_bytes_impl/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro2::Span;
5use proc_macro_hack::proc_macro_hack;
6
7use quote::ToTokens;
8use syn::parse_macro_input;
9
10struct ByteConcat(Vec<u8>);
11
12enum Param {
13    Byte(u8),
14    ByteStr(Vec<u8>),
15}
16
17impl Param {
18    fn len(&self) -> usize {
19        match self {
20            Param::Byte(_) => 1,
21            Param::ByteStr(vec) => vec.len(),
22        }
23    }
24}
25
26impl syn::parse::Parse for ByteConcat {
27    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
28        let mut args = input.parse_terminated::<_, syn::Token![,]>(|input| {
29            if input.peek(syn::LitByte) {
30                Ok(Param::Byte(input.parse::<syn::LitByte>()?.value()))
31            } else {
32                Ok(Param::ByteStr(input.parse::<syn::LitByteStr>()?.value()))
33            }
34        })?;
35
36        let mut bytes = Vec::with_capacity(args.iter().map(|slice| slice.len()).sum());
37
38        for param in args.iter_mut() {
39            match param {
40                Param::Byte(b) => bytes.push(*b),
41                Param::ByteStr(ref mut vec) => bytes.append(vec),
42            }
43        }
44
45        Ok(ByteConcat(bytes))
46    }
47}
48
49#[proc_macro_hack]
50pub fn concat_bytes(input: TokenStream) -> TokenStream {
51    let bytes = parse_macro_input!(input as ByteConcat);
52
53    let literal = syn::LitByteStr::new(&bytes.0, Span::call_site());
54    literal.into_token_stream().into()
55}