lang: Make InitSpace support unnamed and unit structs (#3084)

coral-xyz · Jul 14, 2024 · a2bd50b · a2bd50b
1 parent ca4822f
commit a2bd50b
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 18 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -24,6 +24,7 @@ The minor version will be incremented upon a breaking change and the patch versi
 - lang: Fix `align` repr support in `declare-program!` ([#3056](https://github.com/coral-xyz/anchor/pull/3056)).
 - lang: Make stack frames slimmer on ATA creation ([#3065](https://github.com/coral-xyz/anchor/pull/3065)).
 - lang: Remove `getrandom` dependency ([#3072](https://github.com/coral-xyz/anchor/pull/3072)).
+- lang: Make `InitSpace` support unnamed & unit structs ([#3084](https://github.com/coral-xyz/anchor/pull/3084)).
 
 ### Breaking
 

diff --git a/lang/derive/space/src/lib.rs b/lang/derive/space/src/lib.rs
@@ -4,8 +4,8 @@ use proc_macro::TokenStream;
 use proc_macro2::{Ident, TokenStream as TokenStream2, TokenTree};
 use quote::{quote, quote_spanned, ToTokens};
 use syn::{
-    parse::ParseStream, parse2, parse_macro_input, Attribute, DeriveInput, Fields, GenericArgument,
-    LitInt, PathArguments, Type, TypeArray,
+    parse::ParseStream, parse2, parse_macro_input, punctuated::Punctuated, token::Comma, Attribute,
+    DeriveInput, Field, Fields, GenericArgument, LitInt, PathArguments, Type, TypeArray,
 };
 
 /// Implements a [`Space`](./trait.Space.html) trait on the given
@@ -41,22 +41,30 @@ pub fn derive_init_space(item: TokenStream) -> TokenStream {
     let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
     let name = input.ident;
 
+    let process_struct_fields = |fields: Punctuated<Field, Comma>| {
+        let recurse = fields.into_iter().map(|f| {
+            let mut max_len_args = get_max_len_args(&f.attrs);
+            len_from_type(f.ty, &mut max_len_args)
+        });
+
+        quote! {
+            #[automatically_derived]
+            impl #impl_generics anchor_lang::Space for #name #ty_generics #where_clause {
+                const INIT_SPACE: usize = 0 #(+ #recurse)*;
+            }
+        }
+    };
+
     let expanded: TokenStream2 = match input.data {
         syn::Data::Struct(strct) => match strct.fields {
-            Fields::Named(named) => {
-                let recurse = named.named.into_iter().map(|f| {
-                    let mut max_len_args = get_max_len_args(&f.attrs);
-                    len_from_type(f.ty, &mut max_len_args)
-                });
-
-                quote! {
-                    #[automatically_derived]
-                    impl #impl_generics anchor_lang::Space for #name #ty_generics #where_clause {
-                        const INIT_SPACE: usize = 0 #(+ #recurse)*;
-                    }
+            Fields::Named(named) => process_struct_fields(named.named),
+            Fields::Unnamed(unnamed) => process_struct_fields(unnamed.unnamed),
+            Fields::Unit => quote! {
+                #[automatically_derived]
+                impl #impl_generics anchor_lang::Space for #name #ty_generics #where_clause {
+                    const INIT_SPACE: usize = 0;
                 }
-            }
-            _ => panic!("Please use named fields in account structure"),
+            },
         },
         syn::Data::Enum(enm) => {
             let variants = enm.variants.into_iter().map(|v| {

diff --git a/lang/tests/space.rs b/lang/tests/space.rs
@@ -43,7 +43,7 @@ pub struct TestBasicVarAccount {
 
 #[account]
 #[derive(InitSpace)]
-pub struct TestComplexeVarAccount {
+pub struct TestComplexVarAccount {
     pub test_key: Pubkey,
     #[max_len(10)]
     pub test_vec: Vec<u8>,
@@ -97,6 +97,18 @@ pub struct TestConst {
     pub test_array: [u8; MAX_LEN as usize],
 }
 
+#[derive(InitSpace)]
+pub struct TestUnnamedStruct(
+    pub u8,
+    #[max_len(4)] pub Vec<u32>,
+    #[max_len(10)] pub String,
+    pub ChildStruct,
+    pub TestBasicEnum,
+);
+
+#[derive(InitSpace)]
+pub struct TestUnitStruct;
+
 #[test]
 fn test_empty_struct() {
     assert_eq!(TestEmptyAccount::INIT_SPACE, 0);
@@ -108,9 +120,9 @@ fn test_basic_struct() {
 }
 
 #[test]
-fn test_complexe_struct() {
+fn test_complex_struct() {
     assert_eq!(
-        TestComplexeVarAccount::INIT_SPACE,
+        TestComplexVarAccount::INIT_SPACE,
         32 + 4 + 10 + (4 + 10) + 3
     )
 }
@@ -147,3 +159,16 @@ fn test_full_path() {
 fn test_const() {
     assert_eq!(TestConst::INIT_SPACE, (4 + 10) + 10)
 }
+
+#[test]
+fn test_unnamed_struct() {
+    assert_eq!(
+        TestUnnamedStruct::INIT_SPACE,
+        1 + 4 + 4 * 4 + 4 + 10 + ChildStruct::INIT_SPACE + TestBasicEnum::INIT_SPACE
+    )
+}
+
+#[test]
+fn test_unit_struct() {
+    assert_eq!(TestUnitStruct::INIT_SPACE, 0)
+}