use crate::bpf_writer::BpfWriter;
use crate::error::{Error, ErrorCode};
use crate::{
Accounts, AccountsClose, AccountsExit, Key, Owner, Result, ToAccountInfo, ToAccountInfos,
ToAccountMetas, ZeroCopy,
};
use arrayref::array_ref;
use solana_program::account_info::AccountInfo;
use solana_program::instruction::AccountMeta;
use solana_program::pubkey::Pubkey;
use std::cell::{Ref, RefMut};
use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::io::Write;
use std::marker::PhantomData;
use std::mem;
use std::ops::DerefMut;
#[derive(Clone)]
pub struct AccountLoader<'info, T: ZeroCopy + Owner> {
acc_info: AccountInfo<'info>,
phantom: PhantomData<&'info T>,
}
impl<'info, T: ZeroCopy + Owner + fmt::Debug> fmt::Debug for AccountLoader<'info, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("AccountLoader")
.field("acc_info", &self.acc_info)
.field("phantom", &self.phantom)
.finish()
}
}
impl<'info, T: ZeroCopy + Owner> AccountLoader<'info, T> {
fn new(acc_info: AccountInfo<'info>) -> AccountLoader<'info, T> {
Self {
acc_info,
phantom: PhantomData,
}
}
#[inline(never)]
pub fn try_from(acc_info: &AccountInfo<'info>) -> Result<AccountLoader<'info, T>> {
if acc_info.owner != &T::owner() {
return Err(Error::from(ErrorCode::AccountOwnedByWrongProgram)
.with_pubkeys((*acc_info.owner, T::owner())));
}
let data: &[u8] = &acc_info.try_borrow_data()?;
if data.len() < T::discriminator().len() {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}
let disc_bytes = array_ref![data, 0, 8];
if disc_bytes != &T::discriminator() {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}
Ok(AccountLoader::new(acc_info.clone()))
}
#[inline(never)]
pub fn try_from_unchecked(
_program_id: &Pubkey,
acc_info: &AccountInfo<'info>,
) -> Result<AccountLoader<'info, T>> {
if acc_info.owner != &T::owner() {
return Err(Error::from(ErrorCode::AccountOwnedByWrongProgram)
.with_pubkeys((*acc_info.owner, T::owner())));
}
Ok(AccountLoader::new(acc_info.clone()))
}
pub fn load(&self) -> Result<Ref<T>> {
let data = self.acc_info.try_borrow_data()?;
if data.len() < T::discriminator().len() {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}
let disc_bytes = array_ref![data, 0, 8];
if disc_bytes != &T::discriminator() {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}
Ok(Ref::map(data, |data| {
bytemuck::from_bytes(&data[8..mem::size_of::<T>() + 8])
}))
}
pub fn load_mut(&self) -> Result<RefMut<T>> {
if !self.acc_info.is_writable {
return Err(ErrorCode::AccountNotMutable.into());
}
let data = self.acc_info.try_borrow_mut_data()?;
if data.len() < T::discriminator().len() {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}
let disc_bytes = array_ref![data, 0, 8];
if disc_bytes != &T::discriminator() {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}
Ok(RefMut::map(data, |data| {
bytemuck::from_bytes_mut(&mut data.deref_mut()[8..mem::size_of::<T>() + 8])
}))
}
pub fn load_init(&self) -> Result<RefMut<T>> {
if !self.acc_info.is_writable {
return Err(ErrorCode::AccountNotMutable.into());
}
let data = self.acc_info.try_borrow_mut_data()?;
let mut disc_bytes = [0u8; 8];
disc_bytes.copy_from_slice(&data[..8]);
let discriminator = u64::from_le_bytes(disc_bytes);
if discriminator != 0 {
return Err(ErrorCode::AccountDiscriminatorAlreadySet.into());
}
Ok(RefMut::map(data, |data| {
bytemuck::from_bytes_mut(&mut data.deref_mut()[8..mem::size_of::<T>() + 8])
}))
}
}
impl<'info, T: ZeroCopy + Owner> Accounts<'info> for AccountLoader<'info, T> {
#[inline(never)]
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 Err(ErrorCode::AccountNotEnoughKeys.into());
}
let account = &accounts[0];
*accounts = &accounts[1..];
let l = AccountLoader::try_from(account)?;
Ok(l)
}
}
impl<'info, T: ZeroCopy + Owner> AccountsExit<'info> for AccountLoader<'info, T> {
fn exit(&self, _program_id: &Pubkey) -> Result<()> {
let mut data = self.acc_info.try_borrow_mut_data()?;
let dst: &mut [u8] = &mut data;
let mut writer = BpfWriter::new(dst);
writer.write_all(&T::discriminator()).unwrap();
Ok(())
}
}
impl<'info, T: ZeroCopy + Owner> AccountsClose<'info> for AccountLoader<'info, T> {
fn close(&self, sol_destination: AccountInfo<'info>) -> Result<()> {
crate::common::close(self.to_account_info(), sol_destination)
}
}
impl<'info, T: ZeroCopy + Owner> ToAccountMetas for AccountLoader<'info, T> {
fn to_account_metas(&self, is_signer: Option<bool>) -> Vec<AccountMeta> {
let is_signer = is_signer.unwrap_or(self.acc_info.is_signer);
let meta = match self.acc_info.is_writable {
false => AccountMeta::new_readonly(*self.acc_info.key, is_signer),
true => AccountMeta::new(*self.acc_info.key, is_signer),
};
vec![meta]
}
}
impl<'info, T: ZeroCopy + Owner> AsRef<AccountInfo<'info>> for AccountLoader<'info, T> {
fn as_ref(&self) -> &AccountInfo<'info> {
&self.acc_info
}
}
impl<'info, T: ZeroCopy + Owner> ToAccountInfos<'info> for AccountLoader<'info, T> {
fn to_account_infos(&self) -> Vec<AccountInfo<'info>> {
vec![self.acc_info.clone()]
}
}
impl<'info, T: ZeroCopy + Owner> Key for AccountLoader<'info, T> {
fn key(&self) -> Pubkey {
*self.acc_info.key
}
}