Make stack frames slimmer on ATA creation

coral-xyz · Jul 2, 2024 · 1ae4f51 · 1ae4f51
1 parent e68b96f
commit 1ae4f51
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 12 deletions.
diff --git a/lang/syn/src/codegen/accounts/constraints.rs b/lang/syn/src/codegen/accounts/constraints.rs
@@ -624,17 +624,15 @@ fn generate_constraint_init_group(
                     if !#if_needed || owner_program == &anchor_lang::solana_program::system_program::ID {
                         #payer_optional_check
 
-                        let cpi_program = associated_token_program.to_account_info();
-                        let cpi_accounts = ::anchor_spl::associated_token::Create {
-                            payer: #payer.to_account_info(),
-                            associated_token: #field.to_account_info(),
-                            authority: #owner.to_account_info(),
-                            mint: #mint.to_account_info(),
-                            system_program: system_program.to_account_info(),
-                            token_program: #token_program.to_account_info(),
-                        };
-                        let cpi_ctx = anchor_lang::context::CpiContext::new(cpi_program, cpi_accounts);
-                        ::anchor_spl::associated_token::create(cpi_ctx)?;
+                        ::anchor_spl::associated_token::create_slim_stack_frame(
+                            &#payer.to_account_info(),
+                            &#field.to_account_info(),
+                            &#owner.to_account_info(),
+                            &#mint.to_account_info(),
+                            &system_program.to_account_info(),
+                            &#token_program.to_account_info(),
+                            &#associated_token_program.to_account_info()
+                        )?;
                     }
                     let pa: #ty_decl = #from_account_info_unchecked;
                     if #if_needed {

diff --git a/spl/src/associated_token.rs b/spl/src/associated_token.rs
@@ -8,6 +8,28 @@ pub use spl_associated_token_account::{
     get_associated_token_address, get_associated_token_address_with_program_id, ID,
 };
 
+pub fn create_slim_stack_frame<'info>(
+    payer: &AccountInfo<'info>,
+    associated_token: &AccountInfo<'info>,
+    authority: &AccountInfo<'info>,
+    mint: &AccountInfo<'info>,
+    system_program: &AccountInfo<'info>,
+    token_program: &AccountInfo<'info>,
+    associated_token_program: &AccountInfo<'info>,
+) -> Result<()> {
+    create(CpiContext::new(
+        associated_token_program.clone(),
+        Create {
+            payer: payer.clone(),
+            associated_token: associated_token.clone(),
+            authority: authority.clone(),
+            mint: mint.clone(),
+            system_program: system_program.clone(),
+            token_program: token_program.clone(),
+        },
+    ))
+}
+
 pub fn create<'info>(ctx: CpiContext<'_, '_, '_, 'info, Create<'info>>) -> Result<()> {
     let ix = spl_associated_token_account::instruction::create_associated_token_account(
         ctx.accounts.payer.key,

diff --git a/tests/pda-derivation/programs/pda-derivation/src/lib.rs b/tests/pda-derivation/programs/pda-derivation/src/lib.rs
@@ -85,6 +85,7 @@ pub struct InitMyAccount<'info> {
     base: Account<'info, BaseAccount>,
     // Intentionally using this qualified form instead of importing to test parsing
     another_base: Account<'info, crate::other::AnotherBaseAccount>,
+    /// CHECK:
     base2: AccountInfo<'info>,
     #[account(
         init,
@@ -163,6 +164,43 @@ pub struct AssociatedTokenResolution<'info> {
         associated_token::mint = mint,
     )]
     pub ata: Account<'info, TokenAccount>,
+    #[account(
+        init,
+        payer = payer,
+        associated_token::authority = system_program,
+        associated_token::mint = mint,
+    )]
+    pub ata2: Account<'info, TokenAccount>,
+    #[account(
+        init,
+        payer = payer,
+        associated_token::authority = token_program,
+        associated_token::mint = mint,
+    )]
+    pub ata3: Account<'info, TokenAccount>,
+    #[account(
+        init,
+        payer = payer,
+        associated_token::authority = associated_token_program,
+        associated_token::mint = mint,
+    )]
+    pub ata4: Account<'info, TokenAccount>,
+    #[account(
+        init,
+        payer = payer,
+        associated_token::authority = mint,
+        associated_token::mint = mint,
+    )]
+    pub ata5: Account<'info, TokenAccount>,
+    /// CHECK:
+    pub user: AccountInfo<'info>,
+    #[account(
+        init,
+        payer = payer,
+        associated_token::authority = user,
+        associated_token::mint = mint,
+    )]
+    pub ata6: Account<'info, TokenAccount>,
     #[account(mut)]
     pub payer: Signer<'info>,
     pub system_program: Program<'info, System>,
@@ -174,6 +212,7 @@ pub struct AssociatedTokenResolution<'info> {
 pub struct SeedMathExpr<'info> {
     #[account(seeds = [b"const"], bump)]
     pub my_account: Account<'info, MyAccount>,
+    /// CHECK:
     #[account(seeds = [&(my_account.data + 1).to_le_bytes()], bump)]
     pub math_expr_account: UncheckedAccount<'info>,
 }

diff --git a/tests/pda-derivation/tests/typescript.spec.ts b/tests/pda-derivation/tests/typescript.spec.ts
@@ -110,9 +110,10 @@ describe("typescript", () => {
 
   it("Can resolve associated token accounts", async () => {
     const mintKp = anchor.web3.Keypair.generate();
+    const userKp = anchor.web3.Keypair.generate();
     await program.methods
       .associatedTokenResolution()
-      .accounts({ mint: mintKp.publicKey })
+      .accounts({ mint: mintKp.publicKey, user: userKp.publicKey })
       .signers([mintKp])
       .rpc();
   });