subrpcer_impl/
lib.rs

1// proc-macro
2use proc_macro::TokenStream;
3// crates.io
4use proc_macro2::TokenStream as TokenStream2;
5use syn::{
6	punctuated::Punctuated, spanned::Spanned, token::Comma, Expr, ExprPath, FnArg, Ident, ItemFn,
7	Pat, Path, PathArguments, PathSegment, Stmt,
8};
9
10#[proc_macro_attribute]
11pub fn rpc(_: TokenStream, input: TokenStream) -> TokenStream {
12	let call = syn::parse_macro_input!(input as ItemFn);
13
14	// dbg!(&call);
15
16	let call_vis = call.vis;
17	let call_sig = call.sig;
18	let call_name = call_sig.ident.clone();
19	let call_with_id_name = format!("{}_with_id", call_name.to_string())
20		.parse::<TokenStream2>()
21		.unwrap();
22	let inputs = call_sig.inputs.clone();
23	let input_names = inputs
24		.iter()
25		.map(|a| {
26			if let FnArg::Typed(t) = a {
27				if let Pat::Ident(i) = &*t.pat {
28					i.ident.clone()
29				} else {
30					unimplemented!()
31				}
32			} else {
33				unimplemented!()
34			}
35		})
36		.collect::<Punctuated<_, Comma>>();
37	let output = call_sig.output.clone();
38	let (call_with_id_block, default_id) = {
39		let mut call_with_id_block = call.block;
40		let mut default_id = "DEFAULT_ID".parse::<TokenStream2>().unwrap();
41
42		if let Stmt::Expr(Expr::Call(c)) = &mut call_with_id_block.stmts[0] {
43			if let Expr::Path(p) = &mut c.args[0] {
44				let i = &mut p.path.segments[0].ident;
45
46				default_id = i.to_string().parse().unwrap();
47				*i = Ident::new("id", i.span());
48			}
49		}
50
51		(call_with_id_block, default_id)
52	};
53	let call_with_raw_params_and_id_block = {
54		let mut call_with_raw_params_and_id_block = call_with_id_block.clone();
55
56		if let Stmt::Expr(Expr::Call(c)) = &mut call_with_raw_params_and_id_block.stmts[0] {
57			if let Some(m) = c.args.pop() {
58				c.args.push(Expr::Path(ExprPath {
59					attrs: vec![],
60					qself: None,
61					path: Path {
62						leading_colon: None,
63						segments: vec![PathSegment {
64							ident: Ident::new("params", m.span()),
65							arguments: PathArguments::None,
66						}]
67						.into_iter()
68						.collect(),
69					},
70				}));
71			}
72		}
73
74		call_with_raw_params_and_id_block
75	};
76	let mut token_stream = quote::quote! {
77		#call_vis #call_sig {
78			#call_with_id_name(#default_id, #input_names)
79		}
80		#call_vis fn #call_with_id_name(id: impl Serialize, #inputs) #output #call_with_id_block
81	};
82
83	if !inputs.is_empty() {
84		let call_with_raw_params_name = format!("{}_with_raw_params", call_name.to_string())
85			.parse::<TokenStream2>()
86			.unwrap();
87		let call_with_raw_params_and_id_name =
88			format!("{}_and_id", call_with_raw_params_name.to_string())
89				.parse::<TokenStream2>()
90				.unwrap();
91
92		token_stream.extend(quote::quote! {
93			#[cfg(feature = "raw-params")]
94			#call_vis fn #call_with_raw_params_and_id_name(id: impl Serialize, params: impl Serialize) #output
95			#call_with_raw_params_and_id_block
96			#[cfg(feature = "raw-params")]
97			#call_vis fn #call_with_raw_params_name(params: impl Serialize) #output {
98				#call_with_raw_params_and_id_name(#default_id, params)
99			}
100		});
101	}
102
103	token_stream.into()
104}