Merge pull request #87 from staratlasmeta/stegaBOB/fix/args

stegaBOB · web-flow · commit 0b21f27b5073 · 2024-12-04T09:28:37.000-06:00
Fix: account set args
diff --git a/framework/star_frame/src/account_set/data_account.rs b/framework/star_frame/src/account_set/data_account.rs
@@ -124,6 +124,7 @@ where
     T: ProgramAccount + UnsizedType + ?Sized,
 {
     /// Validates the owner and the discriminant of the account.
+    #[inline]
     pub fn validate(&self) -> Result<()> {
         if self.owner() != &T::OwnerProgram::PROGRAM_ID {
             bail!(ProgramError::IllegalOwner);
@@ -133,6 +134,7 @@ where
         Ok(())
     }
 
+    #[inline]
     fn check_discriminant(bytes: &[u8]) -> Result<()> {
         if bytes.len() < size_of::<<T::OwnerProgram as StarFrameProgram>::AccountDiscriminant>()
             || from_bytes::<PackedValue<<T::OwnerProgram as StarFrameProgram>::AccountDiscriminant>>(
@@ -144,6 +146,7 @@ where
         Ok(())
     }
 
+    #[inline]
     pub fn data<'a>(&'a self) -> Result<RefWrapper<AccountInfoRef<'a>, T::RefData>> {
         let r: Ref<'a, _> = self.info_data_bytes()?;
         Self::check_discriminant(&r)?;
@@ -158,6 +161,7 @@ where
         T::from_bytes(account_info_ref).map(|ret| ret.ref_wrapper)
     }
 
+    #[inline]
     pub fn data_mut<'a>(
         &'a mut self,
     ) -> Result<RefWrapper<AccountInfoRefMut<'a, 'info, T::OwnerProgram>, T::RefData>> {
diff --git a/framework/star_frame/src/util.rs b/framework/star_frame/src/util.rs
@@ -8,6 +8,7 @@ use std::fmt::Debug;
 use std::mem::size_of;
 
 /// Similar to [`Ref::map`], but the closure can return an error.
+#[inline]
 pub fn try_map_ref<'a, I: 'a + ?Sized, O: 'a + ?Sized, E>(
     r: Ref<'a, I>,
     f: impl FnOnce(&I) -> Result<&O, E>,
@@ -20,7 +21,8 @@ pub fn try_map_ref<'a, I: 'a + ?Sized, O: 'a + ?Sized, E>(
     }
 }
 
-/// Similar to [`RefMut::map`], but the closure can return an error.
+/// Similar to [`RefMut::map`], but the closure can return an error
+#[inline]
 pub fn try_map_ref_mut<'a, I: 'a + ?Sized, O: 'a + ?Sized, E>(
     mut r: RefMut<'a, I>,
     f: impl FnOnce(&mut I) -> Result<&mut O, E>,
@@ -60,6 +62,7 @@ unsafe impl<S> RefBytes<S> for OffsetRef
 where
     S: AsBytes,
 {
+    #[inline]
     fn bytes(wrapper: &RefWrapper<S, Self>) -> Result<&[u8]> {
         let mut bytes = wrapper.sup().as_bytes()?;
         bytes.try_advance(wrapper.r().0)?;
@@ -70,6 +73,7 @@ unsafe impl<S> RefBytesMut<S> for OffsetRef
 where
     S: AsMutBytes,
 {
+    #[inline]
     fn bytes_mut(wrapper: &mut RefWrapper<S, Self>) -> Result<&mut [u8]> {
         let (sup, r) = unsafe { wrapper.s_r_mut() };
         let mut bytes = sup.as_mut_bytes()?;
@@ -79,6 +83,7 @@ where
 }
 
 /// Returns a slice of bytes from an array of [`NoUninit`] types.
+#[inline]
 pub fn uninit_array_bytes<T: NoUninit, const N: usize>(array: &[T; N]) -> &[u8] {
     // Safety: `T` is `NoUninit`, so all underlying reads are valid since there's no padding
     // between array elements. The pointer is valid. The entire memory is valid.
diff --git a/framework/star_frame_proc/src/account_set/struct_impl/cleanup.rs b/framework/star_frame_proc/src/account_set/struct_impl/cleanup.rs
@@ -143,7 +143,12 @@ pub(super) fn cleanups(
                     arg: #cleanup_type,
                     syscalls: &mut impl #syscall_invoke<#info_lifetime>,
                 ) -> #result<()> {
-                    #(<#field_type as #account_set_cleanup<#info_lifetime, _>>::cleanup_accounts(&mut self.#field_name, #cleanup_args, syscalls)?;)*
+                    #(
+                        {
+                            let __cleanup_arg = #cleanup_args;
+                            <#field_type as #account_set_cleanup<#info_lifetime, _>>::cleanup_accounts(&mut self.#field_name, __cleanup_arg, syscalls)?;
+                        }
+                    )*
                     #extra_cleanup
                     Ok(())
                 }
diff --git a/framework/star_frame_proc/src/account_set/struct_impl/decode.rs b/framework/star_frame_proc/src/account_set/struct_impl/decode.rs
@@ -165,7 +165,10 @@ pub(super) fn decodes(
         let decode_inner = init(&mut decode_field_ty.iter().zip_eq(&decode_args).map(|(field_ty, decode_args)| {
             match &field_ty {
                 DecodeFieldTy::Type(field_type) => quote! {
-                        <#field_type as #account_set_decode<#decode_lifetime, #info_lifetime, _>>::decode_accounts(accounts, #decode_args, syscalls)?
+                        {
+                            let __arg = #decode_args;
+                            <#field_type as #account_set_decode<#decode_lifetime, #info_lifetime, _>>::decode_accounts(accounts, __arg, syscalls)?
+                        }
                     },
                 DecodeFieldTy::Default(default) => quote!(#default)
             }
diff --git a/framework/star_frame_proc/src/account_set/struct_impl/validate.rs b/framework/star_frame_proc/src/account_set/struct_impl/validate.rs
@@ -182,11 +182,15 @@ pub(super) fn validates(
                 quote! {}
             } else {
                 let address_check = validate_address.as_ref().map(|address| quote! {
-                    <#field_type as #prelude::SingleAccountSet<#info_lifetime>>::check_key(&self.#field_name, #address)?;
+                    let __address = #address;
+                    <#field_type as #prelude::SingleAccountSet<#info_lifetime>>::check_key(&self.#field_name, __address)?;
                 });
                 quote! {
-                    #address_check
-                    <#field_type as #account_set_validate<#info_lifetime, #validate_ty>>::validate_accounts(&mut self.#field_name, #validate_arg, syscalls)?;
+                    {
+                        #address_check
+                        let __validate_arg = #validate_arg;
+                        <#field_type as #account_set_validate<#info_lifetime, #validate_ty>>::validate_accounts(&mut self.#field_name, __validate_arg, syscalls)?;
+                    }
                 }
             })
             .collect::<Vec<_>>();
diff --git a/framework/star_frame_spl/src/token/state.rs b/framework/star_frame_spl/src/token/state.rs
@@ -5,6 +5,7 @@ use star_frame::account_set::AccountSet;
 use star_frame::anyhow::{bail, Context};
 use star_frame::bytemuck;
 use star_frame::prelude::*;
+use star_frame::util::try_map_ref;
 use std::cell::Ref;
 
 /// A wrapper around `AccountInfo` for the [`spl_token::state::Mint`] account.
@@ -51,6 +52,7 @@ impl<'info> MintAccount<'info> {
     /// ```
     pub const LEN: usize = 82;
 
+    #[inline]
     pub fn validate(&mut self) -> Result<()> {
         if self.validated {
             return Ok(());
@@ -62,20 +64,27 @@ impl<'info> MintAccount<'info> {
         if self.info_data_bytes()?.len() != Self::LEN {
             bail!(ProgramError::InvalidAccountData);
         }
+        // set validate before checking state to allow us to call .data()
+        self.validated = true;
         // check initialized
         if !self.data()?.is_initialized {
             bail!(ProgramError::UninitializedAccount);
         }
-        self.validated = true;
         Ok(())
     }
 
+    #[inline]
     pub fn data(&self) -> Result<Ref<MintData>> {
-        Ok(Ref::map(self.info_data_bytes()?, |data| {
-            bytemuck::checked::from_bytes::<MintData>(data)
-        }))
+        if !self.validated {
+            return Err(ProgramError::InvalidAccountData)
+                .context("Called `.data()` on MintAccount before validation");
+        }
+        Ok(try_map_ref(self.info_data_bytes()?, |data| {
+            bytemuck::checked::try_from_bytes::<MintData>(data)
+        })?)
     }
 
+    #[inline]
     pub fn validate_mint(&self, validate_mint: ValidateMint) -> Result<()> {
         let data = self.data()?;
         if let Some(decimals) = validate_mint.decimals {
@@ -266,6 +275,7 @@ impl<'info> TokenAccount<'info> {
     /// ```
     pub const LEN: usize = 165;
 
+    #[inline]
     pub fn validate(&mut self) -> Result<()> {
         if self.validated {
             return Ok(());
@@ -277,19 +287,26 @@ impl<'info> TokenAccount<'info> {
         if self.info_data_bytes()?.len() != Self::LEN {
             bail!(ProgramError::InvalidAccountData);
         }
+        // set validate before checking state to allow us to call .data()
+        self.validated = true;
         if self.data()?.state == AccountState::Uninitialized {
             bail!(ProgramError::UninitializedAccount);
         }
-        self.validated = true;
         Ok(())
     }
 
+    #[inline]
     pub fn data(&self) -> Result<Ref<TokenAccountData>> {
-        Ok(Ref::map(self.info_data_bytes()?, |data| {
-            bytemuck::checked::from_bytes::<TokenAccountData>(data)
-        }))
+        if !self.validated {
+            return Err(ProgramError::InvalidAccountData)
+                .context("Called `.data()` on TokenAccount before validation");
+        }
+        Ok(try_map_ref(self.info_data_bytes()?, |data| {
+            bytemuck::checked::try_from_bytes::<TokenAccountData>(data)
+        })?)
     }
 
+    #[inline]
     pub fn validate_token(&self, validate_token: ValidateToken) -> Result<()> {
         let data = self.data()?;
         if let Some(mint) = validate_token.mint {
