anchor_derive_space/
lib.rs

1use {
2    proc_macro::TokenStream,
3    proc_macro2::{Ident, TokenStream as TokenStream2},
4    quote::{quote, quote_spanned, ToTokens},
5    std::collections::VecDeque,
6    syn::{
7        parse::{Parse, ParseStream},
8        parse_macro_input,
9        punctuated::Punctuated,
10        spanned::Spanned,
11        token::Comma,
12        Attribute, DeriveInput, Expr, ExprLit, Field, Fields, GenericArgument, Lit, PathArguments,
13        Token, Type, TypeArray,
14    },
15};
16
17/// Implements a [`Space`](./trait.Space.html) trait on the given
18/// struct or enum.
19///
20/// For types that have a variable size like String and Vec, it is necessary to indicate the size by the `max_len` attribute.
21/// For nested types, it is necessary to specify a size for each variable type (see example).
22///
23/// # Example
24/// ```ignore
25/// use anchor_lang::prelude::*;
26///
27/// #[account]
28/// #[derive(InitSpace)]
29/// pub struct ExampleAccount {
30///     pub data: u64,
31///     #[max_len(50)]
32///     pub string_one: String,
33///     #[max_len(10, 5)]
34///     pub nested: Vec<Vec<u8>>,
35/// }
36///
37/// #[derive(Accounts)]
38/// pub struct Initialize<'info> {
39///    #[account(mut)]
40///    pub payer: Signer<'info>,
41///    pub system_program: Program<'info, System>,
42///    #[account(init, payer = payer, space = 8 + ExampleAccount::INIT_SPACE)]
43///    pub data: Account<'info, ExampleAccount>,
44/// }
45/// ```
46#[proc_macro_derive(InitSpace, attributes(max_len))]
47pub fn derive_init_space(item: TokenStream) -> TokenStream {
48    let input = parse_macro_input!(item as DeriveInput);
49    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
50    let name = input.ident.clone();
51
52    let process_struct_fields =
53        |fields: Punctuated<Field, Comma>| -> Result<TokenStream2, syn::Error> {
54            let recurse = fields
55                .into_iter()
56                .map(|f| {
57                    let mut max_len_args = get_max_len_args(&f.attrs)?;
58                    Ok(len_from_type(f.ty, &mut max_len_args))
59                })
60                .collect::<Result<Vec<_>, syn::Error>>()?;
61
62            Ok(quote! {
63                #[automatically_derived]
64                impl #impl_generics anchor_lang::Space for #name #ty_generics #where_clause {
65                    const INIT_SPACE: usize = 0 #(+ #recurse)*;
66                }
67            })
68        };
69
70    let expanded = (|| -> Result<TokenStream2, syn::Error> {
71        match input.data {
72            syn::Data::Struct(strct) => match strct.fields {
73                Fields::Named(named) => process_struct_fields(named.named),
74                Fields::Unnamed(unnamed) => process_struct_fields(unnamed.unnamed),
75                Fields::Unit => Ok(quote! {
76                    #[automatically_derived]
77                    impl #impl_generics anchor_lang::Space for #name #ty_generics #where_clause {
78                        const INIT_SPACE: usize = 0;
79                    }
80                }),
81            },
82            syn::Data::Enum(enm) => {
83                let variants = enm
84                    .variants
85                    .into_iter()
86                    .map(|v| {
87                        let len = v
88                            .fields
89                            .into_iter()
90                            .map(|f| {
91                                let mut max_len_args = get_max_len_args(&f.attrs)?;
92                                Ok(len_from_type(f.ty, &mut max_len_args))
93                            })
94                            .collect::<Result<Vec<_>, syn::Error>>()?;
95
96                        Ok(quote! {
97                            0 #(+ #len)*
98                        })
99                    })
100                    .collect::<Result<Vec<_>, syn::Error>>()?;
101
102                let max = gen_max(variants.into_iter());
103
104                Ok(quote! {
105                    #[automatically_derived]
106                    impl anchor_lang::Space for #name {
107                        const INIT_SPACE: usize = 1 + #max;
108                    }
109                })
110            }
111            _ => Err(syn::Error::new(
112                input.ident.span(),
113                "#[derive(InitSpace)] is only supported on structs and enums",
114            )),
115        }
116    })();
117
118    TokenStream::from(match expanded {
119        Ok(expanded) => expanded,
120        Err(err) => err.into_compile_error(),
121    })
122}
123
124fn gen_max<T: Iterator<Item = TokenStream2>>(mut iter: T) -> TokenStream2 {
125    if let Some(item) = iter.next() {
126        let next_item = gen_max(iter);
127        quote!(anchor_lang::__private::max(#item, #next_item))
128    } else {
129        quote!(0)
130    }
131}
132
133fn len_from_type(ty: Type, attrs: &mut Option<VecDeque<TokenStream2>>) -> TokenStream2 {
134    match ty {
135        Type::Array(TypeArray { elem, len, .. }) => {
136            let array_len = len.to_token_stream();
137            let type_len = len_from_type(*elem, attrs);
138            quote!((#array_len * #type_len))
139        }
140        Type::Path(ty_path) => {
141            let path_segment = match ty_path.path.segments.last() {
142                Some(seg) => seg,
143                None => {
144                    return syn::Error::new_spanned(ty_path, "expected a valid type path")
145                        .into_compile_error()
146                }
147            };
148            let ident = &path_segment.ident;
149            let type_name = ident.to_string();
150            let first_ty = get_first_ty_arg(&path_segment.arguments);
151
152            match type_name.as_str() {
153                "i8" | "u8" | "bool" => quote!(1),
154                "i16" | "u16" => quote!(2),
155                "i32" | "u32" | "f32" => quote!(4),
156                "i64" | "u64" | "f64" => quote!(8),
157                "i128" | "u128" => quote!(16),
158                "String" => {
159                    let max_len = get_next_arg(ident, attrs);
160                    quote!((4 + #max_len))
161                }
162                "Pubkey" => quote!(32),
163                "Option" => {
164                    if let Some(ty) = first_ty {
165                        let type_len = len_from_type(ty, attrs);
166
167                        quote!((1 + #type_len))
168                    } else {
169                        quote_spanned!(ident.span() => compile_error!("Invalid argument in Option"))
170                    }
171                }
172                "Vec" => {
173                    if let Some(ty) = first_ty {
174                        let max_len = get_next_arg(ident, attrs);
175                        let type_len = len_from_type(ty, attrs);
176
177                        quote!((4 + #type_len * #max_len))
178                    } else {
179                        quote_spanned!(ident.span() => compile_error!("Invalid argument in Vec"))
180                    }
181                }
182                _ => {
183                    let ty = &ty_path.path;
184                    quote!(<#ty as anchor_lang::Space>::INIT_SPACE)
185                }
186            }
187        }
188        Type::Tuple(ty_tuple) => {
189            let recurse = ty_tuple
190                .elems
191                .iter()
192                .map(|t| len_from_type(t.clone(), attrs));
193            quote! {
194                (0 #(+ #recurse)*)
195            }
196        }
197        _ => {
198            let ty_type = ty.to_token_stream();
199            syn::Error::new_spanned(ty_type, "Type is not supported by `#[derive(InitSpace)]`")
200                .into_compile_error()
201        }
202    }
203}
204
205fn get_first_ty_arg(args: &PathArguments) -> Option<Type> {
206    match args {
207        PathArguments::AngleBracketed(bracket) => bracket.args.iter().find_map(|el| match el {
208            GenericArgument::Type(ty) => Some(ty.to_owned()),
209            _ => None,
210        }),
211        _ => None,
212    }
213}
214
215fn parse_len_arg(item: ParseStream) -> Result<VecDeque<TokenStream2>, syn::Error> {
216    // Parse comma-separated expressions
217    let exprs = item.parse_terminated(Expr::parse, Token![,])?;
218    let mut result = VecDeque::new();
219
220    // Push them in reverse because get_next_arg() pops from the back
221    for expr in exprs.into_iter().rev() {
222        match expr {
223            Expr::Path(path) => result.push_back(quote!((#path as usize))),
224            Expr::Lit(ExprLit {
225                lit: Lit::Int(lit_int),
226                ..
227            }) => result.push_back(quote!(#lit_int as usize)),
228            other => {
229                return Err(syn::Error::new(
230                    other.span(),
231                    "max_len only accepts integer literals, identifiers, or paths",
232                ))
233            }
234        }
235    }
236
237    Ok(result)
238}
239
240fn get_max_len_args(
241    attributes: &[Attribute],
242) -> Result<Option<VecDeque<TokenStream2>>, syn::Error> {
243    attributes
244        .iter()
245        .find(|a| a.path().is_ident("max_len"))
246        .map(|a| a.parse_args_with(parse_len_arg))
247        .transpose()
248}
249
250fn get_next_arg(ident: &Ident, args: &mut Option<VecDeque<TokenStream2>>) -> TokenStream2 {
251    if let Some(arg_list) = args {
252        if let Some(arg) = arg_list.pop_back() {
253            quote!(#arg)
254        } else {
255            quote_spanned!(ident.span() => compile_error!("The number of lengths are invalid."))
256        }
257    } else {
258        quote_spanned!(ident.span() => compile_error!("Expected max_len attribute."))
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use {super::*, syn::parse::Parser};
265
266    #[test]
267    fn parse_len_arg_accepts_int_literals_and_paths() {
268        let mut args = parse_len_arg.parse_str("10, module::MAX_LEN").unwrap();
269
270        assert_eq!(args.pop_back().unwrap().to_string(), "10 as usize");
271        assert_eq!(
272            args.pop_back().unwrap().to_string(),
273            "(module :: MAX_LEN as usize)"
274        );
275    }
276
277    #[test]
278    fn parse_len_arg_rejects_non_integer_literals() {
279        let err = parse_len_arg.parse_str("1.5").unwrap_err();
280
281        assert_eq!(
282            err.to_string(),
283            "max_len only accepts integer literals, identifiers, or paths"
284        );
285    }
286
287    #[test]
288    fn get_max_len_args_propagates_parse_errors() {
289        let attr = syn::parse_quote!(#[max_len(1.5)]);
290
291        let err = get_max_len_args(&[attr]).unwrap_err();
292
293        assert_eq!(
294            err.to_string(),
295            "max_len only accepts integer literals, identifiers, or paths"
296        );
297    }
298}
anchor_derive_space/lib.rs

anchor_derive_space/
lib.rs
