proc_concat_bytes_impl/
lib.rs1extern 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}