lang: Support for const in the InitSpace macro (coral-xyz#2555)

Aursen · Jul 5, 2023 · 9ff7dfc · 9ff7dfc
1 parent 401d526
commit 9ff7dfc
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 13 deletions.
diff --git 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
@@ -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<T: Iterator<Item = TokenStream2>>(mut iter: T) -> TokenStream2 {
     }
 }
 
-fn len_from_type(ty: Type, attrs: &mut Option<IntoIter<LitInt>>) -> TokenStream2 {
+fn len_from_type(ty: Type, attrs: &mut Option<VecDeque<TokenStream2>>) -> 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<Type> {
     }
 }
 
-fn get_max_len_args(attributes: &[Attribute]) -> Option<IntoIter<LitInt>> {
+fn parse_len_arg(item: ParseStream) -> Result<VecDeque<TokenStream2>, 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::<LitInt>(lit.into_token_stream()) {
+                    result.push_front(quote!(#lit_int))
+                }
+            }
+            _ => (),
+        }
+    }
+
+    Ok(result)
+}
+
+fn get_max_len_args(attributes: &[Attribute]) -> Option<VecDeque<TokenStream2>> {
     attributes
         .iter()
         .find(|a| a.path.is_ident("max_len"))
-        .and_then(|a| {
-            a.parse_args_with(Punctuated::<LitInt, Token![,]>::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<IntoIter<LitInt>>) -> TokenStream2 {
+fn get_next_arg(ident: &Ident, args: &mut Option<VecDeque<TokenStream2>>) -> 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
@@ -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)
+}