anchor_modular_program/
lib.rs

1#![warn(missing_docs)]
2
3//! A replacement #\[program\] macro for anchor-lang that allows splitting
4//! instructions into modules.
5//!
6//!
7//! ```
8//! mod extra;                 
9//! use extra::types::*;       
10//!                            
11//! #[modular_program(         
12//!     modules=[              
13//!         extra::instructions
14//!     ]                      
15//! )]                         
16//! mod my_program {           
17//!     use super::*;          
18//! }
19//! ```
20
21
22
23use std::{collections::HashMap, io::Read};
24use anchor_syn::*;
25use syn::*;
26use quote::*;
27
28
29/// Modules can either be a rust path to an instructions module,
30/// or it can be an object:
31///
32/// ```
33/// #[modular_program(modules=[
34///     mymod::instructions,
35///     {
36///         // Required, module path to instructions
37///         module: path::to::instructions,
38///
39///         // Optional path, override where to look for the instructions
40///         file_path: "./src/mod/etc.rs",
41///
42///         // Optional prefix, empty for no prefix
43///         prefix: "prefix",
44///
45///         // Optional, A macro that wraps the call to the instruction, eg:
46///         // ```
47///         // macro_rules ix_wrapper {
48///         //     ($ix:path, $ctx:ident: $ctx_type:ty $(, $arg:ident: $arg_type:ty )*) => {
49///         //         $ix($ctx, $(, $arg)*)
50///         //     };
51///         // }
52///         // ```
53///         macro: path::to::macro
54///     }
55/// ])]
56/// mod my_program {           
57///     use super::*;          
58/// }
59/// ```
60///
61#[proc_macro_attribute]
62pub fn modular_program(
63    args: proc_macro::TokenStream,
64    input: proc_macro::TokenStream,
65) -> proc_macro::TokenStream {
66    let ProgramMacro { modules } = syn::parse_macro_input!(args as ProgramMacro);
67
68    let fns = modules.into_iter()
69        .map(|m| (m.clone(), get_program(m.clone())))
70        .flat_map(|(spec, p)| {
71            p.ixs.into_iter().map(|ix| build_relay(&spec, ix)).collect::<Vec<_>>()
72        })
73        .collect();
74
75    let input = insert_fns_into_first_module(input, fns);
76    let program = syn::parse_macro_input!(input as anchor_syn::Program);
77    program.to_token_stream().into()
78}
79
80
81/*
82 * This builds relay instructions, i.e, for `foo::instructions::do_thing`:
83 *
84 * pub fn foo_do_thing(ctx: Context<YourInstructionContext>, ...) -> Result<()> {
85 *     foo::instructions::do_thing(ctx, ...)
86 * }
87 */
88
89fn build_relay(spec: &ModuleSpec, ix: Ix) -> ItemFn {
90
91    let item_fn = &ix.raw_method;
92    let ItemFn { attrs, .. } = &item_fn;
93    let Signature { ident: fn_name, inputs, generics, output, .. } = &item_fn.sig;
94
95    let path = spec.module.clone();
96    let first = path.segments[0].clone().ident;
97
98    let new_name = match &spec.prefix {
99        Some(s) if s.is_empty() => fn_name.clone(),
100        o => {
101            let prefix = o.clone().unwrap_or(first.to_string());
102            Ident::new(format!("{}_{}", prefix, fn_name).as_str(), first.span())
103        }
104    };
105
106    // Extract argument names for the function call
107    let arg_names: Vec<Ident> = inputs
108        .iter()
109        .filter_map(|arg| match arg { FnArg::Typed(pt) => Some(&*pt.pat), _ => None })
110        .filter_map(|pt| match pt { syn::Pat::Ident(id) => Some(id.ident.clone()), _ => None })
111        .collect();
112
113    if let Some(w) = &spec.wrapper {
114        parse_quote! {
115            #(#attrs)*
116            pub fn #new_name #generics(#inputs) #output {
117                { #w!(#path::#fn_name, #inputs) }
118            }
119        }
120    } else {
121        parse_quote! {
122            #(#attrs)*
123            pub fn #new_name #generics(#inputs) #output {
124                #path::#fn_name(#(#arg_names),*)
125            }
126        }
127    }
128}
129
130
131/*
132 * Get an anchor Program from the given path, by parsing the file, i.e.
133 * foo::instructions is converted to "$PROGRAM_DIR/src/foo/instructions.rs"
134 */
135
136fn get_program(spec: ModuleSpec) -> Program {
137
138    let mod_path = format!(
139        "{}/{}",
140        std::env::var("CARGO_MANIFEST_DIR").unwrap(),
141        spec.get_file_path()
142    );
143
144    let mut code_str = String::new();
145    std::fs::File::open(mod_path).unwrap().read_to_string(&mut code_str).unwrap();
146    let parsed = syn::parse_file(&code_str).unwrap();
147
148    let program_mod = ItemMod {
149        vis: Visibility::Public(VisPublic { pub_token: Default::default() }),
150        attrs: vec![],
151        mod_token: syn::token::Mod::default(),
152        ident: Ident::new("abc", proc_macro2::Span::call_site()),
153        content: Some((
154            syn::token::Brace::default(),
155            parsed.items,
156        )),
157        semi: None,
158    };
159
160    let program = anchor_syn::parser::program::parse(program_mod).unwrap();
161    assert!(program.fallback_fn.is_none(), "additional program module cant have fallback");
162    program
163}
164
165
166
167/*
168 * Parse the macro arguments
169 */
170
171#[derive(Debug)]
172struct ProgramMacro { modules: Vec<ModuleSpec>, }
173
174impl parse::Parse for ProgramMacro {
175    fn parse(input: parse::ParseStream) -> syn::Result<Self> {
176
177        // Parse `modules`
178        let modules_ident: Ident = input.parse()?;
179        if modules_ident != "modules" {
180            return Err(syn::Error::new(modules_ident.span(), "expected `modules`"));
181        }
182
183        input.parse::<Token![=]>()?;
184
185        // Parse the bracketed list `[cell, placement]`
186        let content;
187        syn::bracketed!(content in input);
188        let specs = content.parse_terminated::<ModuleSpec, Token![,]>(|p| p.parse())?;
189
190        // Convert Punctuated<Ident, _> to Vec<Ident>
191        let modules = specs.into_iter().collect();
192
193        Ok(ProgramMacro { modules })
194    }
195}
196
197#[derive(Clone, Debug)]
198struct ModuleSpec {
199    module: Path,
200    prefix: Option<String>,
201    file_path: Option<String>,
202    wrapper: Option<Path>
203}
204
205impl ModuleSpec {
206    fn get_file_path(&self) -> String {
207        self.file_path.clone().unwrap_or_else(|| {
208            let p = self.module.segments.iter().fold(String::new(), |s, p| format!("{}/{}", s, p.ident));
209            format!("./src{}.rs", p)
210        })
211    }
212}
213
214impl parse::Parse for ModuleSpec {
215    fn parse(input: parse::ParseStream) -> Result<Self> {
216
217        type T = (String, (Option<String>, Option<Path>));
218        fn parse_field(p: parse::ParseStream) -> syn::Result<T> {
219            let name = p.parse::<Ident>()?.to_string();
220            p.parse::<Token![:]>()?;
221            Ok(
222                if name == "file_path" || name == "prefix" {
223                    (name, (Some(p.parse::<LitStr>()?.value()), None))
224                } else if name == "module" || name == "wrapper" {
225                    (name, (None, Some(p.parse::<Path>()?)))
226                } else {
227                    panic!("Invalid module spec param: {}", name);
228                }
229            )
230        }
231
232
233        if input.peek(Ident) {
234            let module = input.parse::<Path>()?;
235            Ok(ModuleSpec { module, prefix: None, file_path: None, wrapper: None })
236        } else {
237            let content;
238            syn::braced!(content in input);
239            let fields = content.parse_terminated::<T, Token![,]>(parse_field)?;
240            let mut hm = fields.clone().into_iter().collect::<HashMap<String, _>>();
241            assert!(hm.len() == fields.len(), "duplicate field");
242            Ok(ModuleSpec {
243                module: hm.remove("module").expect("module is required").1.unwrap(),
244                prefix: hm.remove("prefix").map(|t| t.0).flatten(),
245                file_path: hm.remove("file_path").map(|t| t.0).flatten(),
246                wrapper: hm.remove("wrapper").map(|t| t.1).flatten(),
247            })
248        }
249    }
250}
251
252
253/*
254 * Append instruction functions to main program module
255 */
256
257fn insert_fns_into_first_module(input: proc_macro::TokenStream, fns: Vec<ItemFn>) -> proc_macro::TokenStream {
258
259    let mut item_mod: ItemMod = parse2(input.into()).expect("Failed to parse main program module");
260
261    item_mod.content
262        .as_mut()
263        .expect("Program module has no body?")
264        .1.extend(fns.into_iter().map(Into::into));
265
266    quote! { #item_mod }.into()
267}
268
anchor_modular_program/lib.rs

anchor_modular_program/
lib.rs
