Skip to main content

bunpack_proc_macros/
lib.rs

1mod parse;
2
3struct UnpackArgs {
4    fmt: String,
5    data: syn::Expr,
6}
7
8impl syn::parse::Parse for UnpackArgs {
9    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
10        let fmt = input.parse::<syn::LitStr>()?.value();
11        _ = input.parse::<syn::Token![,]>()?;
12        let data = input.parse()?;
13
14        Ok(Self { fmt, data })
15    }
16}
17
18struct UnpackReadArgs {
19    reader: syn::Expr,
20    fmt: String,
21}
22
23impl syn::parse::Parse for UnpackReadArgs {
24    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
25        let reader = input.parse()?;
26        _ = input.parse::<syn::Token![,]>()?;
27        let fmt = input.parse::<syn::LitStr>()?.value();
28
29        Ok(Self { reader, fmt })
30    }
31}
32
33struct PackArgs {
34    fmt: String,
35    args: Vec<syn::Expr>,
36}
37
38impl syn::parse::Parse for PackArgs {
39    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
40        let fmt = input.parse::<syn::LitStr>()?.value();
41        _ = input.parse::<syn::Token![,]>()?;
42        let args = input
43            .parse_terminated(syn::Expr::parse, syn::Token![,])?
44            .into_iter()
45            .collect();
46        Ok(Self { fmt, args })
47    }
48}
49
50struct PackWriteArgs {
51    writer: syn::Expr,
52    fmt: String,
53    args: Vec<syn::Expr>,
54}
55
56impl syn::parse::Parse for PackWriteArgs {
57    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
58        let writer = input.parse()?;
59        _ = input.parse::<syn::Token![,]>()?;
60        let fmt = input.parse::<syn::LitStr>()?.value();
61        _ = input.parse::<syn::Token![,]>()?;
62        let args = input
63            .parse_terminated(syn::Expr::parse, syn::Token![,])?
64            .into_iter()
65            .collect();
66        Ok(Self { writer, fmt, args })
67    }
68}
69
70#[proc_macro]
71pub fn pack(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
72    let PackArgs { fmt, args } = syn::parse_macro_input!(input as PackArgs);
73
74    let (little_endian, types) =
75        parse::fmt_parser::fmt_pack(&fmt).unwrap_or_else(|e| panic!("Invalid format string: {e}"));
76
77    if types.len() != args.len() {
78        panic!(
79            "Number of format specifiers ({}) does not match number of arguments ({})",
80            types.len(),
81            args.len()
82        );
83    }
84
85    let arg_names = (0..args.len())
86        .map(|i| syn::Ident::new(&format!("arg{i}"), proc_macro2::Span::call_site()))
87        .collect::<Vec<_>>();
88    let fn_args = arg_names
89        .iter()
90        .zip(types.iter())
91        .map(|(arg, ty)| quote::quote! { #arg: #ty })
92        .collect::<Vec<_>>();
93
94    quote::quote! {{
95        fn pack( #( #fn_args ),* ) -> Vec<u8>
96        {
97            let mut buf = Vec::new();
98            #(
99                if #little_endian {
100                    <#types as ::bunpack::Pack>::pack_le(#arg_names, &mut buf);
101                } else {
102                    <#types as ::bunpack::Pack>::pack_be(#arg_names, &mut buf);
103                }
104            )*
105            buf
106        }
107
108        pack( #( #args ),* )
109    }}
110    .into()
111}
112
113#[proc_macro]
114pub fn pack_write(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
115    let PackWriteArgs { writer, fmt, args } = syn::parse_macro_input!(input as PackWriteArgs);
116
117    let (little_endian, types) =
118        parse::fmt_parser::fmt_pack(&fmt).unwrap_or_else(|e| panic!("Invalid format string: {e}"));
119
120    if types.len() != args.len() {
121        panic!(
122            "Number of format specifiers ({}) does not match number of arguments ({})",
123            types.len(),
124            args.len()
125        );
126    }
127
128    let arg_names = (0..args.len())
129        .map(|i| syn::Ident::new(&format!("arg{i}"), proc_macro2::Span::call_site()))
130        .collect::<Vec<_>>();
131    let fn_args = arg_names
132        .iter()
133        .zip(types.iter())
134        .map(|(arg, ty)| quote::quote! { #arg: #ty })
135        .collect::<Vec<_>>();
136
137    quote::quote! {{
138        fn pack_write<W: std::io::Write>(writer: &mut W, #( #fn_args ),* ) -> std::io::Result<()>
139        {
140            let mut buf = Vec::new();
141            #(
142                if #little_endian {
143                    <#types as ::bunpack::Pack>::pack_le(#arg_names, &mut buf);
144                } else {
145                    <#types as ::bunpack::Pack>::pack_be(#arg_names, &mut buf);
146                }
147            )*
148
149            writer.write_all(&buf)
150        }
151
152        pack_write(#writer, #( #args ),* )
153    }}
154    .into()
155}
156
157#[proc_macro]
158pub fn unpack(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
159    let UnpackArgs { fmt, data } = syn::parse_macro_input!(input as UnpackArgs);
160
161    let (little_endian, types) = parse::fmt_parser::fmt_unpack(&fmt)
162        .unwrap_or_else(|e| panic!("Invalid format string: {e}"));
163
164    quote::quote! {{
165        fn unpack<T: AsRef<[u8]>>(data: T) -> ( #( #types ),* )
166        {
167            let mut data: &[u8] = data.as_ref();
168            (#(
169                if #little_endian {
170                    <#types as ::bunpack::Unpack>::unpack_le(&mut data)
171                } else {
172                    <#types as ::bunpack::Unpack>::unpack_be(&mut data)
173                }
174            ),*)
175        }
176
177        unpack(#data)
178    }}
179    .into()
180}
181
182#[proc_macro]
183pub fn unpack_read(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
184    let UnpackReadArgs { reader, fmt } = syn::parse_macro_input!(input as UnpackReadArgs);
185
186    let (little_endian, types) = parse::fmt_parser::fmt_unpack(&fmt)
187        .unwrap_or_else(|e| panic!("Invalid format string: {e}"));
188
189    quote::quote! {{
190        fn unpack_read<R: std::io::Read>(reader: &mut R) -> std::io::Result<( #( #types ),* )>
191        {
192            Ok((#({
193                let mut buf = [0u8; <#types as ::bunpack::Unpack>::SIZE];
194                reader.read_exact(&mut buf)?;
195                let value = if #little_endian {
196                    <#types as ::bunpack::Unpack>::unpack_le(&mut &buf[..])
197                } else {
198                    <#types as ::bunpack::Unpack>::unpack_be(&mut &buf[..])
199                };
200                value
201            }),*))
202        }
203
204        unpack_read(&mut #reader)
205    }}
206    .into()
207}