1use {
2 crate::{helpers::PathHelper, Encoding},
3 heck::ToSnakeCase,
4 quote::format_ident,
5 syn::{
6 parse::{Parse, Parser},
7 punctuated::Punctuated,
8 visit::Visit,
9 Expr, FnArg, Ident, LitInt, Pat, Token, Type,
10 },
11};
12
13pub struct InstructionReturnData {
14 pub ty: Option<Type>,
15 pub encoding: Encoding,
16}
17
18pub enum InstructionArg {
19 Type { ty: Box<Type>, encoding: Encoding },
20 Context(Ident),
21}
22
23pub struct Instruction {
24 pub name: Ident,
25 pub args: Vec<(Ident, InstructionArg)>,
26 pub return_data: InstructionReturnData,
27}
28
29impl TryFrom<&syn::ItemFn> for Instruction {
30 type Error = syn::Error;
31
32 fn try_from(value: &syn::ItemFn) -> Result<Self, Self::Error> {
33 let return_data = value
34 .sig
35 .output
36 .get_element_with_inner()
37 .and_then(|(_, inner, _)| inner);
38
39 let mut args = Vec::with_capacity(value.sig.inputs.len());
40 for fn_arg in &value.sig.inputs {
41 let FnArg::Typed(pat_ty) = fn_arg else {
42 continue;
43 };
44
45 let Type::Path(ref ty_path) = *pat_ty.ty else {
46 continue;
47 };
48
49 let (name, ty, size) = ty_path
50 .get_element_with_inner()
51 .ok_or(syn::Error::new_spanned(fn_arg, "Invalid FnArg."))?;
52
53 if name == "ProgramIdArg" || name == "Remaining" || name == "AccountIter" {
54 continue;
55 }
56
57 let arg_name = extract_name(&pat_ty.pat)
58 .unwrap_or(format_ident!("{}", name.to_string().to_snake_case()));
59
60 if name == "Arg" || name == "BorshArg" {
61 args.push((
62 arg_name,
63 InstructionArg::Type {
64 ty: Box::new(
65 ty.ok_or(syn::Error::new_spanned(fn_arg, "Invalid argument type."))?,
66 ),
67 encoding: Encoding::Bytemuck,
68 },
69 ));
70 } else if name == "BorshArg" {
71 args.push((
72 arg_name,
73 InstructionArg::Type {
74 ty: Box::new(
75 ty.ok_or(syn::Error::new_spanned(fn_arg, "Invalid argument type."))?,
76 ),
77 encoding: Encoding::Borsh,
78 },
79 ));
80 } else if name == "Array" {
81 let size = size.ok_or(syn::Error::new_spanned(fn_arg, "Invalid Array type."))?;
82 let ty = ty.ok_or(syn::Error::new_spanned(fn_arg, "Invalid argument type."))?;
83 let Type::Path(path) = ty else {
84 return Err(syn::Error::new_spanned(&arg_name, "Invalid ty_path."));
85 };
86 let (name, _, _) = path
87 .get_element_with_inner()
88 .ok_or(syn::Error::new_spanned(&path, "Invalid Array inner type."))?;
89 for i in 0..size {
90 let arg_name = format_ident!("{arg_name}_{i}");
91 args.push((arg_name, InstructionArg::Context(name.clone())));
92 }
93 } else {
94 args.push((arg_name, InstructionArg::Context(name.clone())));
95 }
96 }
97
98 Ok(Instruction {
99 name: value.sig.ident.clone(),
100 args,
101 return_data: InstructionReturnData {
102 ty: return_data,
103 encoding: Encoding::Bytemuck,
104 },
105 })
106 }
107}
108
109fn extract_name(pat: &Pat) -> Option<Ident> {
110 match pat {
111 Pat::Ident(ident) => Some(ident.ident.clone()),
112 Pat::TupleStruct(tuple_struct) => {
113 let pat = tuple_struct.elems.first()?;
114 extract_name(pat)
115 }
116 _ => None,
117 }
118}
119
120#[derive(Default)]
121pub struct InstructionsList(pub Vec<(usize, Ident)>);
122
123struct RouterEntry {
124 discriminator: LitInt,
125 _arrow_eq: Token![=],
126 _arrow_gt: Token![>],
127 handler_name: Ident,
128}
129
130impl Parse for RouterEntry {
131 fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
132 Ok(RouterEntry {
133 discriminator: input.parse()?,
134 _arrow_eq: input.parse()?,
135 _arrow_gt: input.parse()?,
136 handler_name: input.parse()?,
137 })
138 }
139}
140
141impl TryFrom<&syn::ItemConst> for InstructionsList {
142 type Error = syn::Error;
143
144 fn try_from(value: &syn::ItemConst) -> syn::Result<Self> {
145 let Expr::Macro(expr_macro) = value.expr.as_ref() else {
146 return Err(syn::Error::new_spanned(value, "Invalid router type."));
147 };
148
149 let instructions = Punctuated::<RouterEntry, syn::Token![,]>::parse_terminated
150 .parse2(expr_macro.mac.tokens.clone())?;
151 Ok(Self(
152 instructions
153 .iter()
154 .map(|entry| {
155 Ok((
156 entry.discriminator.base10_parse::<usize>()?,
157 entry.handler_name.clone(),
158 ))
159 })
160 .collect::<Result<_, syn::Error>>()?,
161 ))
162 }
163}
164
165impl<'ast> Visit<'ast> for InstructionsList {
166 fn visit_item_const(&mut self, i: &'ast syn::ItemConst) {
167 if i.ident != "ROUTER" {
168 return;
169 }
170
171 if let Ok(ix_list) = InstructionsList::try_from(i) {
172 *self = ix_list;
173 }
174 }
175}
176
177#[cfg(test)]
178mod tests {
179 use {
180 super::*,
181 syn::{parse_quote, ItemConst, ItemFn},
182 };
183
184 #[test]
185 fn test_instruction_list() {
186 let router: ItemConst = parse_quote! {
187 pub const ROUTER: EntryFn = basic_router! {
188 0 => account_iter,
189 1 => initialize,
190 2 => assert
191 };
192 };
193
194 let ix_list = InstructionsList::try_from(&router).unwrap();
195 assert_eq!(ix_list.0[0].0, 0);
196 assert_eq!(ix_list.0[1].0, 1);
197 assert_eq!(ix_list.0[2].0, 2);
198 assert_eq!(ix_list.0[0].1, "account_iter");
199 assert_eq!(ix_list.0[1].1, "initialize");
200 assert_eq!(ix_list.0[2].1, "assert");
201 }
202
203 #[test]
204 fn test_instruction_construction() {
205 let fn_raw: ItemFn = parse_quote! {
206 pub fn instruction_1(ctx: Context1, array: Array<Context2, 2>, arg: Arg<u64>) -> ProgramResult {
207 Ok(())
208 }
209 };
210 let ix = Instruction::try_from(&fn_raw).unwrap();
211
212 assert_eq!(ix.name, "instruction_1");
213 assert_eq!(ix.args.len(), 4);
214 assert_eq!(ix.args[0].0, "ctx");
215 assert!(matches!(&ix.args[0].1, InstructionArg::Context(x) if x == "Context1"));
216 assert_eq!(ix.args[1].0, "array_0");
217 assert!(matches!(&ix.args[1].1, InstructionArg::Context(x) if x == "Context2"));
218 assert_eq!(ix.args[2].0, "array_1");
219 assert!(matches!(&ix.args[2].1, InstructionArg::Context(x) if x == "Context2"));
220 assert_eq!(ix.args[3].0, "arg");
221 assert!(ix.return_data.ty.is_none());
222 assert!(matches!(ix.return_data.encoding, Encoding::Bytemuck));
223 }
224}