lang: Add state size override (#121)

CHANGELOG.md

@@ -14,6 +14,7 @@ incremented for features.
 ## Features
 
 * cli: Specify test files to run ([#118](https://github.com/project-serum/anchor/pull/118)).
+* lang: Allow overriding the `#[state]` account's size ([#121](https://github.com/project-serum/anchor/pull/121)).
 
 ## [0.3.0] - 2021-03-12
 

examples/misc/programs/misc/src/lib.rs

@@ -8,6 +8,20 @@ use anchor_lang::prelude::*;
 #[program]
 pub mod misc {
     use super::*;
+
+    pub const SIZE: u64 = 99;
+
+    #[state(SIZE)]
+    pub struct MyState {
+        pub v: Vec<u8>,
+    }
+
+    impl MyState {
+        pub fn new(_ctx: Context<Ctor>) -> Result<Self, ProgramError> {
+            Ok(Self { v: vec![] })
+        }
+    }
+
     pub fn initialize(ctx: Context<Initialize>, udata: u128, idata: i128) -> ProgramResult {
         ctx.accounts.data.udata = udata;
         ctx.accounts.data.idata = idata;
@@ -15,6 +29,9 @@ pub mod misc {
     }
 }
 
+#[derive(Accounts)]
+pub struct Ctor {}
+
 #[derive(Accounts)]
 pub struct Initialize<'info> {
     #[account(init)]

examples/misc/tests/misc.js

@@ -5,10 +5,19 @@ const assert = require("assert");
 describe("misc", () => {
   // Configure the client to use the local cluster.
   anchor.setProvider(anchor.Provider.env());
+  const program = anchor.workspace.Misc;
+
+  it("Can allocate extra space for a state constructor", async () => {
+    const tx = await program.state.rpc.new();
+    const addr = await program.state.address();
+    const state = await program.state();
+    const accountInfo = await program.provider.connection.getAccountInfo(addr);
+    assert.ok(state.v.equals(Buffer.from([])));
+    assert.ok(accountInfo.data.length === 99);
+  });
 
   it("Can use u128 and i128", async () => {
     const data = new anchor.web3.Account();
-    const program = anchor.workspace.Misc;
     const tx = await program.rpc.initialize(
       new anchor.BN(1234),
       new anchor.BN(22),

lang/attribute/account/src/lib.rs

@@ -19,23 +19,15 @@ use syn::parse_macro_input;
 /// and the account deserialization will exit with an error.
 #[proc_macro_attribute]
 pub fn account(
-    args: proc_macro::TokenStream,
+    _args: proc_macro::TokenStream,
     input: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
-    let namespace = args.to_string().replace("\"", "");
-
     let account_strct = parse_macro_input!(input as syn::ItemStruct);
     let account_name = &account_strct.ident;
 
     let discriminator: proc_macro2::TokenStream = {
         // Namespace the discriminator to prevent collisions.
-        let discriminator_preimage = {
-            if namespace.is_empty() {
-                format!("account:{}", account_name.to_string())
-            } else {
-                format!("{}:{}", namespace, account_name.to_string())
-            }
-        };
+        let discriminator_preimage = format!("account:{}", account_name.to_string());
         let mut discriminator = [0u8; 8];
         discriminator.copy_from_slice(
             &anchor_syn::hash::hash(discriminator_preimage.as_bytes()).to_bytes()[..8],
@@ -57,7 +49,6 @@ pub fn account(
         }
 
         impl anchor_lang::AccountDeserialize for #account_name {
-
             fn try_deserialize(buf: &mut &[u8]) -> std::result::Result<Self, ProgramError> {
                  if buf.len() < #discriminator.len() {
                     return Err(ProgramError::AccountDataTooSmall);

lang/attribute/state/src/lib.rs

@@ -5,15 +5,52 @@ use syn::parse_macro_input;
 
 /// The `#[state]` attribute defines the program's state struct, i.e., the
 /// program's global account singleton giving the program the illusion of state.
+///
+/// To allocate space into the account on initialization, pass in the account
+/// size into the macro, e.g., `#[state(SIZE)]`. Otherwise, the size of the
+/// account returned by the struct's `new` constructor will determine the
+/// account size. When determining a size, make sure to reserve enough space
+/// for the 8 byte account discriminator prepended to the account. That is,
+/// always use 8 extra bytes.
 #[proc_macro_attribute]
 pub fn state(
-    _args: proc_macro::TokenStream,
+    args: proc_macro::TokenStream,
     input: proc_macro::TokenStream,
 ) -> proc_macro::TokenStream {
     let item_struct = parse_macro_input!(input as syn::ItemStruct);
+    let struct_ident = &item_struct.ident;
+
+    let size_override = {
+        if args.is_empty() {
+            // No size override given. The account size is whatever is given
+            // as the initialized value. Use the default implementation.
+            quote! {
+                impl anchor_lang::AccountSize for #struct_ident {
+                    fn size(&self) -> Result<u64, ProgramError> {
+                        Ok(8 + self
+                           .try_to_vec()
+                           .map_err(|_| ProgramError::Custom(1))?
+                           .len() as u64)
+                    }
+                }
+            }
+        } else {
+            let size = proc_macro2::TokenStream::from(args);
+            // Size override given to the macro. Use it.
+            quote! {
+                impl anchor_lang::AccountSize for #struct_ident {
+                    fn size(&self) -> Result<u64, ProgramError> {
+                        Ok(#size)
+                    }
+                }
+            }
+        }
+    };
 
     proc_macro::TokenStream::from(quote! {
         #[account]
         #item_struct
+
+        #size_override
     })
 }

lang/src/lib.rs

@@ -169,6 +169,12 @@ pub trait InstructionData: AnchorSerialize {
     fn data(&self) -> Vec<u8>;
 }
 
+/// Calculates the size of an account, which may be larger than the deserialized
+/// data in it. This trait is currently only used for `#[state]` accounts.
+pub trait AccountSize: AnchorSerialize {
+    fn size(&self) -> Result<u64, ProgramError>;
+}
+
 /// The prelude contains all commonly used components of the crate.
 /// All programs should include it via `anchor_lang::prelude::*;`.
 pub mod prelude {

lang/syn/src/codegen/program.rs

@@ -420,9 +420,8 @@ pub fn generate_non_inlined_handlers(program: &Program) -> proc_macro2::TokenStr
                         let seed = anchor_lang::ProgramState::<#name>::seed();
                         let owner = ctor_accounts.program.key;
                         let to = Pubkey::create_with_seed(&base, seed, owner).unwrap();
-                        // Add 8 for the account discriminator.
-                        let space = 8 + instance.try_to_vec().map_err(|_| ProgramError::Custom(1))?.len();
-                        let lamports = ctor_accounts.rent.minimum_balance(space);
+                        let space = anchor_lang::AccountSize::size(&instance)?;
+                        let lamports = ctor_accounts.rent.minimum_balance(std::convert::TryInto::try_into(space).unwrap());
                         let seeds = &[&[nonce][..]];
                         let ix = anchor_lang::solana_program::system_instruction::create_account_with_seed(
                             from,
@@ -430,7 +429,7 @@ pub fn generate_non_inlined_handlers(program: &Program) -> proc_macro2::TokenStr
                             &base,
                             seed,
                             lamports,
-                            space as u64,
+                            space,
                             owner,
                         );
                         anchor_lang::solana_program::program::invoke_signed(