lang: add the InitSpace macro

lang/Cargo.toml

@@ -32,6 +32,7 @@ anchor-attribute-error = { path = "./attribute/error", version = "0.26.0" }
 anchor-attribute-program = { path = "./attribute/program", version = "0.26.0" }
 anchor-attribute-event = { path = "./attribute/event", version = "0.26.0" }
 anchor-derive-accounts = { path = "./derive/accounts", version = "0.26.0" }
+anchor-derive-space = { path = "./derive/space", version = "0.26.0" }
 arrayref = "0.3.6"
 base64 = "0.13.0"
 borsh = "0.9"

lang/derive/space/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "anchor-derive-space"
+version = "0.26.0"
+authors = ["Serum Foundation <foundation@projectserum.com>"]
+repository = "https://github.com/coral-xyz/anchor"
+license = "Apache-2.0"
+description = "Anchor Derive macro to automatically calculate the size of a structure or an enum"
+rust-version = "1.59"
+edition = "2021"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "1.0"
+quote = "1.0"
+syn = "1.0"

lang/derive/space/src/lib.rs

@@ -0,0 +1,122 @@
+use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::{quote, ToTokens};
+use syn::{
+    parse_macro_input, Attribute, DeriveInput, Fields, GenericArgument, LitInt, PathArguments,
+    Type, TypeArray,
+};
+
+#[proc_macro_derive(InitSpace, attributes(max_len))]
+pub fn derive_anchor_deserialize(item: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(item as DeriveInput);
+    let name = input.ident;
+
+    let expanded: TokenStream2 = match input.data {
+        syn::Data::Struct(strct) => match strct.fields {
+            Fields::Named(named) => {
+                let recurse = named
+                    .named
+                    .into_iter()
+                    .map(|f| len_from_type(f.ty, &f.attrs));
+
+                quote! {
+                    #[automatically_derived]
+                    impl anchor_lang::Space for #name {
+                        const INIT_SPACE: u64 = 0 #(+ #recurse)*;
+                    }
+                }
+            }
+            _ => panic!("Please use named fields in account structure"),
+        },
+        syn::Data::Enum(enm) => {
+            let variants = enm.variants.into_iter().map(|v| {
+                let len = v.fields.into_iter().map(|f| len_from_type(f.ty, &f.attrs));
+
+                quote! {
+                    0 #(+ #len)*
+                }
+            });
+
+            let max = gen_max(variants);
+
+            quote! {
+                #[automatically_derived]
+                impl anchor_lang::Space for #name {
+                    const INIT_SPACE: u64 = 1 + #max;
+                }
+            }
+        }
+        _ => unimplemented!(),
+    };
+
+    TokenStream::from(expanded)
+}
+
+fn gen_max<T: Iterator<Item = TokenStream2>>(mut iter: T) -> TokenStream2 {
+    if let Some(item) = iter.next() {
+        let next_item = gen_max(iter);
+        quote!(anchor_lang::__private::max(#item, #next_item))
+    } else {
+        quote!(0)
+    }
+}
+
+fn len_from_type(ty: Type, attrs: &[Attribute]) -> TokenStream2 {
+    match ty {
+        Type::Array(TypeArray { elem, len, .. }) => {
+            let array_len = len.to_token_stream();
+            let type_len = len_from_type(*elem, attrs);
+            quote!(#array_len * #type_len)
+        }
+        Type::Path(path) => {
+            let path_segment = path.path.segments.last().unwrap();
+            let type_name = path_segment.ident.to_string();
+
+            match type_name.as_str() {
+                "i8" | "u8" | "bool" => quote!(1),
+                "i16" | "u16" => quote!(2),
+                "i32" | "u32" | "f32" => quote!(4),
+                "i64" | "u64" | "f64" => quote!(8),
+                "i128" | "u128" => quote!(16),
+                "String" => {
+                    let max_len = get_max_len(attrs);
+                    quote!(4 + #max_len )
+                }
+                "Pubkey" => quote!(32),
+                "Vec" => match &path_segment.arguments {
+                    PathArguments::AngleBracketed(args) => {
+                        let ty = args
+                            .args
+                            .iter()
+                            .find_map(|el| match el {
+                                GenericArgument::Type(ty) => Some(ty.to_owned()),
+                                _ => None,
+                            })
+                            .unwrap();
+                        let max_len = get_max_len(attrs);
+                        let type_len = len_from_type(ty, attrs);
+
+                        quote!(4 + (#type_len * #max_len))
+                    }
+                    _ => panic!("Invalid argument in Vec"),
+                },
+                _ => {
+                    let ty = &path_segment.ident;
+                    quote!(<#ty as anchor_lang::Space>::INIT_SPACE)
+                }
+            }
+        }
+        _ => panic!("Type {:?} is not supported", ty),
+    }
+}
+
+fn get_max_len(attributes: &[Attribute]) -> u64 {
+    let attr = attributes
+        .iter()
+        .find(|a| a.path.is_ident("max_len"))
+        .expect("Expected max_len attribute");
+
+    attr.parse_args::<LitInt>()
+        .and_then(|l| l.base10_parse())
+        .expect("Can't parse the max_len value")
+}

lang/src/lib.rs

@@ -51,6 +51,7 @@ pub use anchor_attribute_error::*;
 pub use anchor_attribute_event::{emit, event};
 pub use anchor_attribute_program::program;
 pub use anchor_derive_accounts::Accounts;
+pub use anchor_derive_space::InitSpace;
 /// Borsh is the default serialization format for instructions and accounts.
 pub use borsh::{BorshDeserialize as AnchorDeserialize, BorshSerialize as AnchorSerialize};
 pub use solana_program;
@@ -209,6 +210,11 @@ pub trait Discriminator {
     }
 }
 
+/// Defines the space of an account for initialization.
+pub trait Space {
+    const INIT_SPACE: u64;
+}
+
 /// Bump seed for program derived addresses.
 pub trait Bump {
     fn seed(&self) -> u8;
@@ -247,7 +253,7 @@ pub mod prelude {
         require, require_eq, require_gt, require_gte, require_keys_eq, require_keys_neq,
         require_neq, solana_program::bpf_loader_upgradeable::UpgradeableLoaderState, source,
         system_program::System, zero_copy, AccountDeserialize, AccountSerialize, Accounts,
-        AccountsClose, AccountsExit, AnchorDeserialize, AnchorSerialize, Id, Key, Owner,
+        AccountsClose, AccountsExit, AnchorDeserialize, AnchorSerialize, Id, InitSpace, Key, Owner,
         ProgramData, Result, ToAccountInfo, ToAccountInfos, ToAccountMetas,
     };
     pub use anchor_attribute_error::*;
@@ -288,6 +294,13 @@ pub mod __private {
 
     use solana_program::pubkey::Pubkey;
 
+    // Used to calculate the maximum between two expressions.
+    // It is necessary for the calculation of the enum space.
+    #[doc(hidden)]
+    pub const fn max(a: u64, b: u64) -> u64 {
+        [a, b][(a < b) as usize]
+    }
+
     // Very experimental trait.
     #[doc(hidden)]
     pub trait ZeroCopyAccessor<Ty> {

lang/tests/space.rs

@@ -0,0 +1,103 @@
+use anchor_lang::{prelude::*, Space};
+
+// Needed to declare accounts.
+declare_id!("Fg6PaFpoGXkYsidMpWTK6W2BeZ7FEfcYkg476zPFsLnS");
+
+#[derive(InitSpace)]
+pub enum TestBasicEnum {
+    Basic1,
+    Basic2 {
+        test_u8: u8,
+    },
+    Basic3 {
+        test_u16: u16,
+    },
+    Basic4 {
+        #[max_len(10)]
+        test_vec: Vec<u8>,
+    },
+}
+
+#[account]
+#[derive(InitSpace)]
+pub struct TestEmptyAccount {}
+
+#[account]
+#[derive(InitSpace)]
+pub struct TestBasicVarAccount {
+    pub test_u8: u8,
+    pub test_u16: u16,
+    pub test_u32: u32,
+    pub test_u64: u64,
+    pub test_u128: u128,
+}
+
+#[account]
+#[derive(InitSpace)]
+pub struct TestComplexeVarAccount {
+    pub test_key: Pubkey,
+    #[max_len(10)]
+    pub test_vec: Vec<u8>,
+    #[max_len(10)]
+    pub test_string: String,
+}
+
+#[derive(InitSpace)]
+pub struct TestNonAccountStruct {
+    pub test_bool: bool,
+}
+
+#[account(zero_copy)]
+#[derive(InitSpace)]
+pub struct TestZeroCopyStruct {
+    pub test_array: [u8; 10],
+    pub test_u32: u32,
+}
+
+#[derive(InitSpace)]
+pub struct ChildStruct {
+    #[max_len(10)]
+    pub test_string: String,
+}
+
+#[derive(InitSpace)]
+pub struct TestNestedStruct {
+    pub test_struct: ChildStruct,
+    pub test_enum: TestBasicEnum,
+}
+
+#[test]
+fn test_empty_struct() {
+    assert_eq!(TestEmptyAccount::INIT_SPACE, 0);
+}
+
+#[test]
+fn test_basic_struct() {
+    assert_eq!(TestBasicVarAccount::INIT_SPACE, 1 + 2 + 4 + 8 + 16);
+}
+
+#[test]
+fn test_complexe_struct() {
+    assert_eq!(
+        TestComplexeVarAccount::INIT_SPACE,
+        32 + (4 + (1 * 10)) + (4 + 10)
+    )
+}
+
+#[test]
+fn test_zero_copy_struct() {
+    assert_eq!(TestZeroCopyStruct::INIT_SPACE, (1 * 10) + 4)
+}
+
+#[test]
+fn test_basic_enum() {
+    assert_eq!(TestBasicEnum::INIT_SPACE, 1 + 14);
+}
+
+#[test]
+fn test_nested_struct() {
+    assert_eq!(
+        TestNestedStruct::INIT_SPACE,
+        (4 + 10) + TestBasicEnum::INIT_SPACE
+    )
+}