diff --git a/CHANGELOG.md b/CHANGELOG.md index 205f7a4cb8..3d1d7f9922 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ The minor version will be incremented upon a breaking change and the patch versi ### Fixes - ts: Packages no longer depend on `assert` ([#2535](https://github.com/coral-xyz/anchor/pull/2535)). +- lang: Support for `const` in the `InitSpace` macro ([#2555](https://github.com/coral-xyz/anchor/pull/2555)). ### Breaking diff --git a/lang/derive/space/src/lib.rs b/lang/derive/space/src/lib.rs index bd045b0c4e..f60565a9df 100644 --- a/lang/derive/space/src/lib.rs +++ b/lang/derive/space/src/lib.rs @@ -1,10 +1,11 @@ +use std::collections::VecDeque; + use proc_macro::TokenStream; -use proc_macro2::{Ident, TokenStream as TokenStream2}; +use proc_macro2::{Ident, TokenStream as TokenStream2, TokenTree}; use quote::{quote, quote_spanned, ToTokens}; use syn::{ - parse_macro_input, - punctuated::{IntoIter, Punctuated}, - Attribute, DeriveInput, Fields, GenericArgument, LitInt, PathArguments, Token, Type, TypeArray, + parse::ParseStream, parse2, parse_macro_input, Attribute, DeriveInput, Fields, GenericArgument, + LitInt, PathArguments, Type, TypeArray, }; /// Implements a [`Space`](./trait.Space.html) trait on the given @@ -93,7 +94,7 @@ fn gen_max>(mut iter: T) -> TokenStream2 { } } -fn len_from_type(ty: Type, attrs: &mut Option>) -> TokenStream2 { +fn len_from_type(ty: Type, attrs: &mut Option>) -> TokenStream2 { match ty { Type::Array(TypeArray { elem, len, .. }) => { let array_len = len.to_token_stream(); @@ -156,20 +157,33 @@ fn get_first_ty_arg(args: &PathArguments) -> Option { } } -fn get_max_len_args(attributes: &[Attribute]) -> Option> { +fn parse_len_arg(item: ParseStream) -> Result, syn::Error> { + let mut result = VecDeque::new(); + while let Some(token_tree) = item.parse()? { + match token_tree { + TokenTree::Ident(ident) => result.push_front(quote!((#ident as usize))), + TokenTree::Literal(lit) => { + if let Ok(lit_int) = parse2::(lit.into_token_stream()) { + result.push_front(quote!(#lit_int)) + } + } + _ => (), + } + } + + Ok(result) +} + +fn get_max_len_args(attributes: &[Attribute]) -> Option> { attributes .iter() .find(|a| a.path.is_ident("max_len")) - .and_then(|a| { - a.parse_args_with(Punctuated::::parse_terminated) - .ok() - }) - .map(|p| p.into_iter()) + .and_then(|a| a.parse_args_with(parse_len_arg).ok()) } -fn get_next_arg(ident: &Ident, args: &mut Option>) -> TokenStream2 { +fn get_next_arg(ident: &Ident, args: &mut Option>) -> TokenStream2 { if let Some(arg_list) = args { - if let Some(arg) = arg_list.next() { + if let Some(arg) = arg_list.pop_back() { quote!(#arg) } else { quote_spanned!(ident.span() => compile_error!("The number of lengths are invalid.")) diff --git a/lang/tests/space.rs b/lang/tests/space.rs index 4cf81adceb..74c03321d2 100644 --- a/lang/tests/space.rs +++ b/lang/tests/space.rs @@ -88,6 +88,15 @@ pub struct TestFullPath { pub test_path: inside_mod::Data, } +const MAX_LEN: u8 = 10; + +#[derive(InitSpace)] +pub struct TestConst { + #[max_len(MAX_LEN)] + pub test_string: String, + pub test_array: [u8; MAX_LEN as usize], +} + #[test] fn test_empty_struct() { assert_eq!(TestEmptyAccount::INIT_SPACE, 0); @@ -133,3 +142,8 @@ fn test_matrix_struct() { fn test_full_path() { assert_eq!(TestFullPath::INIT_SPACE, 8 + 9) } + +#[test] +fn test_const() { + assert_eq!(TestConst::INIT_SPACE, (4 + 10) + 10) +}