Stega bob/feat/optional accounts

lang/src/accounts/mod.rs

@@ -13,6 +13,7 @@ pub mod cpi_state;
 #[doc(hidden)]
 #[allow(deprecated)]
 pub mod loader;
+pub mod option;
 pub mod program;
 #[doc(hidden)]
 #[allow(deprecated)]

lang/src/accounts/option.rs

@@ -0,0 +1,95 @@
+//! Option<T> type for optional accounts.
+//!
+//! # Example
+//! ```ignore
+//! #[derive(Accounts)]
+//! pub struct Example {
+//!     pub my_acc: Option<Account<'info, MyData>>
+//! }
+//! ```
+
+use std::collections::{BTreeMap, BTreeSet};
+
+use solana_program::account_info::AccountInfo;
+use solana_program::instruction::AccountMeta;
+use solana_program::pubkey::Pubkey;
+
+use crate::{
+    error::ErrorCode, Accounts, AccountsClose, AccountsExit, Result, ToAccountInfo, ToAccountInfos,
+    ToAccountMetas, ToOptionalAccountInfos, TryKey, TryToAccountInfo,
+};
+
+impl<'info, T: Accounts<'info>> Accounts<'info> for Option<T> {
+    fn try_accounts(
+        program_id: &Pubkey,
+        accounts: &mut &[AccountInfo<'info>],
+        ix_data: &[u8],
+        bumps: &mut BTreeMap<String, u8>,
+        reallocs: &mut BTreeSet<Pubkey>,
+    ) -> Result<Self> {
+        if accounts.is_empty() {
+            return Ok(None);
+        }
+        let account = &accounts[0];
+        if account.key == program_id {
+            *accounts = &accounts[1..];
+            Ok(None)
+        } else {
+            T::try_accounts(program_id, accounts, ix_data, bumps, reallocs).map(Some)
+        }
+    }
+}
+
+impl<'info, T: ToAccountInfos<'info>> ToAccountInfos<'info> for Option<T> {
+    fn to_account_infos(&self) -> Vec<AccountInfo<'info>> {
+        self.as_ref()
+            .expect("Cannot run `to_account_infos` on None")
+            .to_account_infos()
+    }
+}
+
+impl<'info, T: ToAccountInfos<'info>> ToOptionalAccountInfos<'info> for Option<T> {
+    fn to_optional_account_infos(&self, program: &AccountInfo<'info>) -> Vec<AccountInfo<'info>> {
+        match self.as_ref() {
+            None => program.to_account_infos(),
+            Some(account) => account.to_account_infos(),
+        }
+    }
+}
+
+impl<T: ToAccountMetas> ToAccountMetas for Option<T> {
+    fn to_account_metas(&self, is_signer: Option<bool>) -> Vec<AccountMeta> {
+        self.as_ref()
+            .expect("Cannot run `to_account_metas` on None")
+            .to_account_metas(is_signer)
+    }
+}
+
+impl<'info, T: AccountsClose<'info>> AccountsClose<'info> for Option<T> {
+    fn close(&self, sol_destination: AccountInfo<'info>) -> Result<()> {
+        self.as_ref()
+            .map_or(Ok(()), |t| T::close(t, sol_destination))
+    }
+}
+
+impl<'info, T: AccountsExit<'info>> AccountsExit<'info> for Option<T> {
+    fn exit(&self, program_id: &Pubkey) -> Result<()> {
+        self.as_ref().map_or(Ok(()), |t| t.exit(program_id))
+    }
+}
+
+impl<T: TryKey> TryKey for Option<T> {
+    fn try_key(&self) -> Result<Pubkey> {
+        self.as_ref()
+            .map_or(Err(ErrorCode::TryKeyOnNone.into()), |t| t.try_key())
+    }
+}
+
+impl<'info, T: ToAccountInfo<'info>> TryToAccountInfo<'info> for Option<T> {
+    fn try_to_account_info(&self) -> Result<AccountInfo<'info>> {
+        self.as_ref()
+            .map_or(Err(ErrorCode::TryKeyOnNone.into()), |t| {
+                Ok(t.to_account_info())
+            })
+    }
+}

lang/src/error.rs

@@ -184,6 +184,9 @@ pub enum ErrorCode {
     /// 3017 - The account was duplicated for more than one reallocation
     #[msg("The account was duplicated for more than one reallocation")]
     AccountDuplicateReallocs,
+    /// 3018 - Tried to get the key from None
+    #[msg("Tried to get the key from None")]
+    TryKeyOnNone,
 
     // State.
     /// 4000 - The given state account does not have the correct address

lang/src/lib.rs

@@ -121,6 +121,14 @@ pub trait ToAccountInfos<'info> {
     fn to_account_infos(&self) -> Vec<AccountInfo<'info>>;
 }
 
+/// Transformation to
+/// [`AccountInfo`](../solana_program/account_info/struct.AccountInfo.html)
+/// structs. Indended for use for `Accounts` structs with optional accounts in order to
+/// pass in the program `AccountInfo`.
+pub trait ToOptionalAccountInfos<'info> {
+    fn to_optional_account_infos(&self, program: &AccountInfo<'info>) -> Vec<AccountInfo<'info>>;
+}
+
 /// Transformation to an `AccountInfo` struct.
 pub trait ToAccountInfo<'info> {
     fn to_account_info(&self) -> AccountInfo<'info>;
@@ -135,6 +143,10 @@ where
     }
 }
 
+pub trait TryToAccountInfo<'info> {
+    fn try_to_account_info(&self) -> Result<AccountInfo<'info>>;
+}
+
 /// A data structure that can be serialized and stored into account storage,
 /// i.e. an
 /// [`AccountInfo`](../solana_program/account_info/struct.AccountInfo.html#structfield.data)'s
@@ -230,6 +242,17 @@ impl Key for Pubkey {
     }
 }
 
+/// Defines the Pubkey of an account in an operation that may fail
+pub trait TryKey {
+    fn try_key(&self) -> Result<Pubkey>;
+}
+
+impl<T: Key> TryKey for T {
+    fn try_key(&self) -> Result<Pubkey> {
+        Ok(self.key())
+    }
+}
+
 /// The prelude contains all commonly used components of the crate.
 /// All programs should include it via `anchor_lang::prelude::*;`.
 pub mod prelude {
@@ -243,7 +266,8 @@ pub mod prelude {
         require_neq, solana_program::bpf_loader_upgradeable::UpgradeableLoaderState, source, state,
         system_program::System, zero_copy, AccountDeserialize, AccountSerialize, Accounts,
         AccountsExit, AnchorDeserialize, AnchorSerialize, Id, Key, Owner, ProgramData, Result,
-        ToAccountInfo, ToAccountInfos, ToAccountMetas,
+        ToAccountInfo, ToAccountInfos, ToAccountMetas, ToOptionalAccountInfos, TryKey,
+        TryToAccountInfo,
     };
     pub use anchor_attribute_error::*;
     pub use borsh;

lang/syn/src/codegen/accounts/__client_accounts.rs

@@ -61,9 +61,16 @@ pub fn generate(accs: &AccountsStruct) -> proc_macro2::TokenStream {
                 } else {
                     quote!()
                 };
-                quote! {
-                    #docs
-                    pub #name: anchor_lang::solana_program::pubkey::Pubkey
+                if f.is_optional {
+                    quote! {
+                        #docs
+                        pub #name: Option<anchor_lang::solana_program::pubkey::Pubkey>
+                    }
+                } else {
+                    quote! {
+                        #docs
+                        pub #name: anchor_lang::solana_program::pubkey::Pubkey
+                    }
                 }
             }
         })
@@ -93,8 +100,18 @@ pub fn generate(accs: &AccountsStruct) -> proc_macro2::TokenStream {
                     true => quote! { anchor_lang::solana_program::instruction::AccountMeta::new },
                 };
                 let name = &f.ident;
-                quote! {
-                    account_metas.push(#meta(self.#name, #is_signer));
+                if f.is_optional {
+                    quote! {
+                        if let Some(#name) = &self.#name {
+                            account_metas.push(#meta(*#name, #is_signer));
+                        } else {
+                            account_metas.push(#meta(ID, #is_signer));
+                        }
+                    }
+                } else {
+                    quote! {
+                        account_metas.push(#meta(self.#name, #is_signer));
+                    }
                 }
             }
         })

lang/syn/src/codegen/accounts/__cpi_client_accounts.rs

@@ -62,9 +62,16 @@ pub fn generate(accs: &AccountsStruct) -> proc_macro2::TokenStream {
                 } else {
                     quote!()
                 };
-                quote! {
-                    #docs
-                    pub #name: anchor_lang::solana_program::account_info::AccountInfo<'info>
+                if f.is_optional {
+                    quote! {
+                        #docs
+                        pub #name: Option<anchor_lang::solana_program::account_info::AccountInfo<'info>>
+                    }
+                } else {
+                    quote! {
+                        #docs
+                        pub #name: anchor_lang::solana_program::account_info::AccountInfo<'info>
+                    }
                 }
             }
         })
@@ -94,8 +101,18 @@ pub fn generate(accs: &AccountsStruct) -> proc_macro2::TokenStream {
                     true => quote! { anchor_lang::solana_program::instruction::AccountMeta::new },
                 };
                 let name = &f.ident;
-                quote! {
-                    account_metas.push(#meta(anchor_lang::Key::key(&self.#name), #is_signer));
+                if f.is_optional {
+                    quote! {
+                        if let Some(#name) = &self.#name {
+                            account_metas.push(#meta(anchor_lang::Key::key(#name), #is_signer));
+                        } else {
+                            account_metas.push(#meta(ID, #is_signer));
+                        }
+                    }
+                } else {
+                    quote! {
+                        account_metas.push(#meta(anchor_lang::Key::key(&self.#name), #is_signer));
+                    }
                 }
             }
         })
@@ -113,8 +130,17 @@ pub fn generate(accs: &AccountsStruct) -> proc_macro2::TokenStream {
             }
             AccountField::Field(f) => {
                 let name = &f.ident;
-                quote! {
-                    account_infos.push(anchor_lang::ToAccountInfo::to_account_info(&self.#name));
+                //TODO: figure out how to handle None gen
+                if f.is_optional {
+                    quote! {
+                        if let Some(account) = &self.#name {
+                            account_infos.push(anchor_lang::ToAccountInfo::to_account_info(account));
+                        }
+                    }
+                } else {
+                    quote! {
+                        account_infos.push(anchor_lang::ToAccountInfo::to_account_info(&self.#name));
+                    }
                 }
             }
         })