Skip to content

Commit

Permalink
lang: remove redundant discriminator checks
Browse files Browse the repository at this point in the history
  • Loading branch information
wowinter13 committed Jan 16, 2025
1 parent db6da5c commit 904b328
Show file tree
Hide file tree
Showing 5 changed files with 7 additions and 67 deletions.
48 changes: 1 addition & 47 deletions lang/src/accounts/account_loader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ use std::ops::DerefMut;
/// Note that using accounts in this way is distinctly different from using,
/// for example, the [`Account`](crate::accounts::account::Account). Namely,
/// one must call
/// - `load_init` after initializing an account (this will ignore the missing
/// account discriminator that gets added only after the user's instruction code)
/// - `load` when the account is not mutable
/// - `load_mut` when the account is mutable
///
Expand Down Expand Up @@ -51,7 +49,7 @@ use std::ops::DerefMut;
/// use super::*;
///
/// pub fn create_bar(ctx: Context<CreateBar>, data: u64) -> Result<()> {
/// let bar = &mut ctx.accounts.bar.load_init()?;
/// let bar = &mut ctx.accounts.bar.load_mut()?;
/// bar.authority = ctx.accounts.authority.key();
/// bar.data = data;
/// Ok(())
Expand Down Expand Up @@ -154,15 +152,6 @@ impl<'info, T: ZeroCopy + Owner> AccountLoader<'info, T> {
pub fn load(&self) -> Result<Ref<T>> {
let data = self.acc_info.try_borrow_data()?;
let disc = T::DISCRIMINATOR;
if data.len() < disc.len() {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}

let given_disc = &data[..disc.len()];
if given_disc != disc {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}

Ok(Ref::map(data, |data| {
bytemuck::from_bytes(&data[disc.len()..mem::size_of::<T>() + disc.len()])
}))
Expand All @@ -178,41 +167,6 @@ impl<'info, T: ZeroCopy + Owner> AccountLoader<'info, T> {

let data = self.acc_info.try_borrow_mut_data()?;
let disc = T::DISCRIMINATOR;
if data.len() < disc.len() {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}

let given_disc = &data[..disc.len()];
if given_disc != disc {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}

Ok(RefMut::map(data, |data| {
bytemuck::from_bytes_mut(
&mut data.deref_mut()[disc.len()..mem::size_of::<T>() + disc.len()],
)
}))
}

/// Returns a `RefMut` to the account data structure for reading or writing.
/// Should only be called once, when the account is being initialized.
pub fn load_init(&self) -> Result<RefMut<T>> {
// AccountInfo api allows you to borrow mut even if the account isn't
// writable, so add this check for a better dev experience.
if !self.acc_info.is_writable {
return Err(ErrorCode::AccountNotMutable.into());
}

let data = self.acc_info.try_borrow_mut_data()?;

// The discriminator should be zero, since we're initializing.
let disc = T::DISCRIMINATOR;
let given_disc = &data[..disc.len()];
let has_disc = given_disc.iter().any(|b| *b != 0);
if has_disc {
return Err(ErrorCode::AccountDiscriminatorAlreadySet.into());
}

Ok(RefMut::map(data, |data| {
bytemuck::from_bytes_mut(
&mut data.deref_mut()[disc.len()..mem::size_of::<T>() + disc.len()],
Expand Down
2 changes: 1 addition & 1 deletion tests/chat/programs/chat/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ pub mod chat {
let given_name = name.as_bytes();
let mut name = [0u8; 280];
name[..given_name.len()].copy_from_slice(given_name);
let mut chat = ctx.accounts.chat_room.load_init()?;
let mut chat = ctx.accounts.chat_room.load_mut()?;
chat.name = name;
Ok(())
}
Expand Down
11 changes: 2 additions & 9 deletions tests/misc/programs/misc-optional/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -109,13 +109,6 @@ pub mod misc_optional {
Ok(())
}

pub fn test_pda_init_zero_copy(ctx: Context<TestPdaInitZeroCopy>) -> Result<()> {
let mut acc = ctx.accounts.my_pda.as_ref().unwrap().load_init()?;
acc.data = 9;
acc.bump = ctx.bumps.my_pda.unwrap();
Ok(())
}

pub fn test_pda_mut_zero_copy(ctx: Context<TestPdaMutZeroCopy>) -> Result<()> {
let mut acc = ctx.accounts.my_pda.as_mut().unwrap().load_mut()?;
acc.data = 1234;
Expand All @@ -139,8 +132,8 @@ pub mod misc_optional {
Ok(())
}

pub fn test_init_zero_copy(ctx: Context<TestInitZeroCopy>) -> Result<()> {
let mut data = ctx.accounts.data.as_ref().unwrap().load_init()?;
pub fn test_mut_zero_copy(ctx: Context<TestInitZeroCopy>) -> Result<()> {
let mut data = ctx.accounts.data.as_ref().unwrap().load_mut()?;
data.data = 10;
data.bump = 2;
Ok(())
Expand Down
9 changes: 1 addition & 8 deletions tests/misc/programs/misc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -103,13 +103,6 @@ pub mod misc {
Ok(())
}

pub fn test_pda_init_zero_copy(ctx: Context<TestPdaInitZeroCopy>) -> Result<()> {
let mut acc = ctx.accounts.my_pda.load_init()?;
acc.data = 9;
acc.bump = ctx.bumps.my_pda;
Ok(())
}

pub fn test_pda_mut_zero_copy(ctx: Context<TestPdaMutZeroCopy>) -> Result<()> {
let mut acc = ctx.accounts.my_pda.load_mut()?;
acc.data = 1234;
Expand All @@ -134,7 +127,7 @@ pub mod misc {
}

pub fn test_init_zero_copy(ctx: Context<TestInitZeroCopy>) -> Result<()> {
let mut data = ctx.accounts.data.load_init()?;
let mut data = ctx.accounts.data.load_mut()?;
data.data = 10;
data.bump = 2;
Ok(())
Expand Down
4 changes: 2 additions & 2 deletions tests/zero-copy/programs/zero-copy/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ pub mod zero_copy {
use super::*;

pub fn create_foo(ctx: Context<CreateFoo>) -> Result<()> {
let foo = &mut ctx.accounts.foo.load_init()?;
let foo = &mut ctx.accounts.foo.load_mut()?;
foo.authority = *ctx.accounts.authority.key;
foo.set_second_authority(ctx.accounts.authority.key);
Ok(())
Expand All @@ -33,7 +33,7 @@ pub mod zero_copy {
}

pub fn create_bar(ctx: Context<CreateBar>) -> Result<()> {
let bar = &mut ctx.accounts.bar.load_init()?;
let bar = &mut ctx.accounts.bar.load_mut()?;
bar.authority = *ctx.accounts.authority.key;
Ok(())
}
Expand Down

0 comments on commit 904b328

Please sign in to comment.