1use std::str::FromStr;
2
3use convert_case::Casing;
4use paste::paste;
5use proc_macro2::{TokenStream, TokenTree};
6use quote::quote;
7use syn::spanned::Spanned;
8
9struct RecordParams(syn::Type, syn::Ident, Option<syn::LitStr>);
10impl syn::parse::Parse for RecordParams {
11 fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
12 let value_type = input.parse()?;
13 let _ = input.parse::<syn::Token![,]>();
14 let ident = input.parse()?;
15 let _ = input.parse::<syn::Token![,]>();
16 let header = input.parse().ok();
17 Ok(RecordParams(value_type, ident, header))
18 }
19}
20
21const GENERIC_ERROR_MSG: &str = "Expected to be used with an enum";
22
23#[proc_macro_attribute]
24pub fn record(
25 params: proc_macro::TokenStream,
26 ast: proc_macro::TokenStream,
27) -> proc_macro::TokenStream {
28 let RecordParams(value_type, ident, header) = syn::parse2(params.into()).expect("Invalid params");
29
30 let ast_tokens: TokenStream = ast.clone().into();
31
32 let mut ast_iter = ast_tokens.clone().into_iter();
33
34 let _ = ast_iter.next().expect(GENERIC_ERROR_MSG);
35 let _ = ast_iter.next().expect(GENERIC_ERROR_MSG);
36
37 let TokenTree::Group(enum_props_group) = ast_iter.next().expect(GENERIC_ERROR_MSG) else {
38 panic!("{}", GENERIC_ERROR_MSG)
39 };
40
41 let struct_props: Vec<TokenStream> = enum_props_group
42 .stream()
43 .into_iter()
44 .filter_map(|token| {
45 if let TokenTree::Ident(ident) = token {
46 let ident_str = ident.to_string();
47 let snake = ident_str.to_case(convert_case::Case::Snake);
48 let snake_ident = syn::Ident::new(&snake, ident_str.span());
49
50 Some(quote! {
51 #snake_ident: #value_type ,
52 })
53 } else {
54 None
55 }
56 })
57 .collect();
58
59 let header_token = header
60 .map(|token| TokenStream::from_str(&token.value()).expect(GENERIC_ERROR_MSG))
61 .unwrap_or(TokenStream::new());
62
63 let output = quote! {
64 #ast_tokens
65 #header_token
66 struct #ident {
67 #(#struct_props)*
68 }
69 };
70
71 output.into()
72}