kaptn_derive_extrametas/
lib.rs

1use proc_macro::TokenStream;
2use quote::{quote, ToTokens};
3use syn::{
4    parse::{Parse, ParseStream},
5    parse_macro_input,
6    punctuated::Punctuated,
7    Attribute, Data, DeriveInput, Expr, Fields, GenericParam, Lit, Meta, NestedMeta, Token,
8};
9
10struct Seeds(Punctuated<Expr, Token![,]>);
11
12impl Parse for Seeds {
13    fn parse(input: ParseStream) -> syn::Result<Self> {
14        let content;
15        syn::bracketed!(content in input);
16        Ok(Seeds(content.parse_terminated(Expr::parse)?))
17    }
18}
19
20#[proc_macro_derive(ExtraMetas, attributes(meta))]
21pub fn derive_extra_metas(input: TokenStream) -> TokenStream {
22    let input = parse_macro_input!(input as DeriveInput);
23    let name = &input.ident;
24
25    // Extract lifetimes
26    let lifetimes: Vec<_> = input
27        .generics
28        .params
29        .iter()
30        .filter_map(|param| {
31            if let GenericParam::Lifetime(lt) = param {
32                Some(lt.lifetime.clone())
33            } else {
34                None
35            }
36        })
37        .collect();
38
39    // Determine if the 'info lifetime is present
40    let has_info_lifetime = lifetimes.iter().any(|lt| lt.to_string() == "'info");
41
42    // Extract fields
43    let fields = if let Data::Struct(data_struct) = &input.data {
44        if let Fields::Named(fields) = &data_struct.fields {
45            fields.named.iter().collect::<Vec<_>>()
46        } else {
47            vec![]
48        }
49    } else {
50        vec![]
51    };
52
53    // Process account metas: pubkey or seeds
54    let account_metas = fields
55        .iter()
56        .filter_map(|f| {
57            let ident = f.ident.as_ref()?;
58            let meta_attr = f.attrs.iter().find(|attr| attr.path.is_ident("meta"))?;
59            parse_meta_attribute(ident, meta_attr).ok()
60        })
61        .collect::<Vec<_>>();
62
63    let field_names = fields
64        .iter()
65        .filter_map(|f| f.ident.as_ref())
66        .collect::<Vec<_>>();
67
68    // Generate from_accounts method
69    let from_accounts = if has_info_lifetime {
70        quote! {
71            fn from_accounts(accounts: &[AccountInfo<'info>]) -> Result<Self, ProgramError> {
72                let mut iter = accounts.iter().skip(5); // Adjust skip as needed
73                Ok(Self {
74                    #(#field_names: iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?.clone(),)*
75                })
76            }
77        }
78    } else {
79        quote! {
80            fn from_accounts(accounts: &[AccountInfo]) -> Result<Self, ProgramError> {
81                let mut iter = accounts.iter().skip(5); // Skip till the extra account metas
82                Ok(Self {
83                    #(#field_names: iter.next().ok_or(ProgramError::NotEnoughAccountKeys)?.clone(),)*
84                })
85            }
86        }
87    };
88
89    // Generate trait implementation for ExtraMetas
90    let gen = if has_info_lifetime {
91        quote! {
92            impl<#(#lifetimes),*> ExtraMetas<#(#lifetimes),*> for #name<#(#lifetimes),*> {
93                #from_accounts
94
95                fn to_extra_account_metas() -> Vec<ExtraAccountMeta> {
96                    vec![
97                        #(#account_metas),*
98                    ]
99                }
100            }
101        }
102    } else if !lifetimes.is_empty() {
103        quote! {
104            impl<#(#lifetimes),*> ExtraMetas<'_, #(#lifetimes),*> for #name<#(#lifetimes),*> {
105                #from_accounts
106
107                fn to_extra_account_metas() -> Vec<ExtraAccountMeta> {
108                    vec![
109                        #(#account_metas),*
110                    ]
111                }
112            }
113        }
114    } else {
115        quote! {
116            impl ExtraMetas<'_> for #name {
117                #from_accounts
118
119                fn to_extra_account_metas() -> Vec<ExtraAccountMeta> {
120                    vec![
121                        #(#account_metas),*
122                    ]
123                }
124            }
125        }
126    };
127
128    gen.into()
129}
130
131fn parse_meta_attribute(
132    _ident: &syn::Ident,
133    attr: &Attribute,
134) -> syn::Result<proc_macro2::TokenStream> {
135    let meta = attr.parse_meta()?;
136    if let Meta::List(list) = meta {
137        let mut pubkey = None;
138        let mut signer = false;
139        let mut writable = false;
140        let mut seeds = None;
141
142        for nested in list.nested.iter() {
143            match nested {
144                NestedMeta::Meta(Meta::NameValue(nv)) => {
145                    match nv.path.get_ident().map(|i| i.to_string()).as_deref() {
146                        Some("pubkey") => {
147                            if let Lit::Str(lit) = &nv.lit {
148                                pubkey = Some(lit.value());
149                            }
150                        }
151                        Some("signer") => {
152                            if let Lit::Bool(lit) = &nv.lit {
153                                signer = lit.value();
154                            }
155                        }
156                        Some("writable") => {
157                            if let Lit::Bool(lit) = &nv.lit {
158                                writable = lit.value();
159                            }
160                        }
161                        _ => {}
162                    }
163                }
164                NestedMeta::Meta(Meta::List(list)) if list.path.is_ident("seeds") => {
165                    seeds = Some(syn::parse2::<Seeds>(list.nested.to_token_stream())?);
166                }
167                _ => {}
168            }
169        }
170
171        if pubkey.is_some() && seeds.is_some() {
172            return Err(syn::Error::new_spanned(
173                attr,
174                "Cannot specify both pubkey and seeds",
175            ));
176        }
177
178        if pubkey.is_none() && seeds.is_none() {
179            return Err(syn::Error::new_spanned(
180                attr,
181                "Must specify either pubkey or seeds",
182            ));
183        }
184
185        // Generate appropriate ExtraAccountMeta based on pubkey or seeds
186        Ok(if let Some(pk) = pubkey {
187            quote! {
188                ExtraAccountMeta::new_with_pubkey(&#pk.parse().unwrap(), #signer, #writable).unwrap()
189            }
190        } else if let Some(Seeds(seeds)) = seeds {
191            let seed_exprs = seeds.iter();
192            quote! {
193                ExtraAccountMeta::new_external_pda_with_seeds(
194                    0, // Associated token program index, adjust as needed
195                    &[#(#seed_exprs),*],
196                    #signer,
197                    #writable
198                ).unwrap()
199            }
200        } else {
201            unreachable!()
202        })
203    } else {
204        Err(syn::Error::new_spanned(
205            attr,
206            "Expected list-style attribute",
207        ))
208    }
209}
kaptn_derive_extrametas/lib.rs

kaptn_derive_extrametas/
lib.rs
