lib.rs 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. extern crate proc_macro;
  2. use quote::quote;
  3. use syn::parse_macro_input;
  4. /// The `#[state]` attribute defines the program's state struct, i.e., the
  5. /// program's global account singleton giving the program the illusion of state.
  6. ///
  7. /// To allocate space into the account on initialization, pass in the account
  8. /// size into the macro, e.g., `#[state(SIZE)]`. Otherwise, the size of the
  9. /// account returned by the struct's `new` constructor will determine the
  10. /// account size. When determining a size, make sure to reserve enough space
  11. /// for the 8 byte account discriminator prepended to the account. That is,
  12. /// always use 8 extra bytes.
  13. #[proc_macro_attribute]
  14. pub fn state(
  15. args: proc_macro::TokenStream,
  16. input: proc_macro::TokenStream,
  17. ) -> proc_macro::TokenStream {
  18. let item_struct = parse_macro_input!(input as syn::ItemStruct);
  19. let struct_ident = &item_struct.ident;
  20. let size_override = {
  21. if args.is_empty() {
  22. // No size override given. The account size is whatever is given
  23. // as the initialized value. Use the default implementation.
  24. quote! {
  25. impl anchor_lang::AccountSize for #struct_ident {
  26. fn size(&self) -> std::result::Result<u64, anchor_lang::solana_program::program_error::ProgramError> {
  27. Ok(8 + self
  28. .try_to_vec()
  29. .map_err(|_| ProgramError::Custom(1))?
  30. .len() as u64)
  31. }
  32. }
  33. }
  34. } else {
  35. let size = proc_macro2::TokenStream::from(args);
  36. // Size override given to the macro. Use it.
  37. quote! {
  38. impl anchor_lang::AccountSize for #struct_ident {
  39. fn size(&self) -> std::result::Result<u64, anchor_lang::solana_program::program_error::ProgramError> {
  40. Ok(#size)
  41. }
  42. }
  43. }
  44. }
  45. };
  46. proc_macro::TokenStream::from(quote! {
  47. #[account]
  48. #item_struct
  49. #size_override
  50. })
  51. }