From 5a6ad4b1322482604aeacfc162de41b03a221760 Mon Sep 17 00:00:00 2001 From: ananas Date: Wed, 21 Jan 2026 02:35:51 +0000 Subject: [PATCH] fix: move token account creation to preinit --- .../macros/src/light_pdas/accounts/builder.rs | 198 +++++++++++------- .../macros/src/light_pdas/accounts/derive.rs | 77 +++---- .../macros/src/light_pdas/accounts/token.rs | 27 +-- 3 files changed, 180 insertions(+), 122 deletions(-) diff --git a/sdk-libs/macros/src/light_pdas/accounts/builder.rs b/sdk-libs/macros/src/light_pdas/accounts/builder.rs index 4c01f9ebf2..9c5bf9c316 100644 --- a/sdk-libs/macros/src/light_pdas/accounts/builder.rs +++ b/sdk-libs/macros/src/light_pdas/accounts/builder.rs @@ -209,47 +209,142 @@ impl LightAccountsBuilder { /// 2. Invoke CreateMintsCpi with CPI context offset /// After this, Mints are "hot" and usable in instruction body pub fn generate_pre_init_pdas_and_mints(&self) -> Result { + let body = self.generate_pre_init_pdas_and_mints_body()?; + Ok(quote! { + #body + Ok(true) + }) + } + + /// Generate LightPreInit body for mints-only (no PDAs): + /// Invoke CreateMintsCpi with decompress directly + /// After this, Mints are "hot" and usable in instruction body + pub fn generate_pre_init_mints_only(&self) -> Result { + let body = self.generate_pre_init_mints_only_body()?; + Ok(quote! { + #body + Ok(true) + }) + } + + /// Generate LightPreInit body for PDAs only (no mints) + /// After this, compressed addresses are registered + pub fn generate_pre_init_pdas_only(&self) -> Result { + let body = self.generate_pre_init_pdas_only_body()?; + Ok(quote! { + #body + Ok(true) + }) + } + + /// Generate unified pre_init body for ALL account types. + /// + /// This method handles all combinations of: + /// - PDAs (compressed accounts) + /// - Mints (compressed mints) + /// - Token accounts (vaults) + /// - ATAs (associated token accounts) + /// + /// ALL account creation happens here so accounts are available during + /// the instruction handler for transfers, minting, etc. + pub fn generate_pre_init_all(&self) -> Result { + let has_pdas = self.has_pdas(); + let has_mints = self.has_mints(); + + // Generate token/ATA creation code (if any) + let token_creation = TokenAccountsBuilder::new( + &self.parsed.token_account_fields, + &self.parsed.ata_fields, + &self.infra, + ) + .generate_pre_init_token_creation(); + + // Handle different combinations + match (has_pdas, has_mints, token_creation.is_some()) { + // PDAs + Mints + Tokens/ATAs + (true, true, true) => { + let pda_mint_body = self.generate_pre_init_pdas_and_mints_body()?; + let token_body = token_creation.unwrap(); + Ok(quote! { + #pda_mint_body + #token_body + Ok(true) + }) + } + // PDAs + Mints (no tokens) + (true, true, false) => self.generate_pre_init_pdas_and_mints(), + // PDAs + Tokens/ATAs (no mints) + (true, false, true) => { + let pda_body = self.generate_pre_init_pdas_only_body()?; + let token_body = token_creation.unwrap(); + Ok(quote! { + #pda_body + #token_body + Ok(true) + }) + } + // PDAs only + (true, false, false) => self.generate_pre_init_pdas_only(), + // Mints + Tokens/ATAs (no PDAs) + (false, true, true) => { + let mint_body = self.generate_pre_init_mints_only_body()?; + let token_body = token_creation.unwrap(); + Ok(quote! { + #mint_body + #token_body + Ok(true) + }) + } + // Mints only + (false, true, false) => self.generate_pre_init_mints_only(), + // Tokens/ATAs only (no PDAs, no mints) + (false, false, true) => { + let token_body = token_creation.unwrap(); + Ok(quote! { + #token_body + Ok(true) + }) + } + // Nothing to do + (false, false, false) => Ok(quote! { Ok(false) }), + } + } + + /// Generate PDAs + mints body WITHOUT the Ok(true) return. + fn generate_pre_init_pdas_and_mints_body(&self) -> Result { let (compress_blocks, new_addr_idents) = generate_pda_compress_blocks(&self.parsed.rentfree_fields); let rentfree_count = self.parsed.rentfree_fields.len() as u8; let pda_count = self.parsed.rentfree_fields.len(); - // Get instruction param ident let first_arg = self.get_first_instruction_arg()?; let params_ident = &first_arg.name; - // Get the first PDA's output tree index (for the state tree output queue) let first_pda_output_tree = &self.parsed.rentfree_fields[0].output_tree; - // Generate CreateMintsCpi invocation for all mints with PDA context offset let mints = &self.parsed.light_mint_fields; let mint_invocation = LightMintsBuilder::new(mints, params_ident, &self.infra) .with_pda_context(pda_count, quote! { #first_pda_output_tree }) .generate_invocation(); - // Infrastructure field references for quote! interpolation let fee_payer = &self.infra.fee_payer; let compression_config = &self.infra.compression_config; Ok(quote! { - // Build CPI accounts WITH CPI context for batching let cpi_accounts = light_sdk::cpi::v2::CpiAccounts::new_with_config( &self.#fee_payer, _remaining, ::light_sdk::sdk_types::CpiAccountsConfig::new_with_cpi_context(crate::LIGHT_CPI_SIGNER), ); - // Load compression config let compression_config_data = light_sdk::interface::LightConfig::load_checked( &self.#compression_config, &crate::ID )?; - // Collect compressed infos for all rentfree PDA accounts let mut all_compressed_infos = Vec::with_capacity(#rentfree_count as usize); #(#compress_blocks)* - // Step 1: Write PDAs to CPI context let cpi_context_account = cpi_accounts.cpi_context()?; let cpi_context_accounts = ::light_sdk::sdk_types::CpiContextWriteAccounts { fee_payer: cpi_accounts.fee_payer(), @@ -268,78 +363,37 @@ impl LightAccountsBuilder { .write_to_cpi_context_first() .invoke_write_to_cpi_context_first(cpi_context_accounts)?; - // Step 2: Create mints using CreateMintsCpi with CPI context offset #mint_invocation - - Ok(true) }) } - /// Generate LightPreInit body for mints-only (no PDAs): - /// Invoke CreateMintsCpi with decompress directly - /// After this, Mints are "hot" and usable in instruction body - pub fn generate_pre_init_mints_only(&self) -> Result { - // Get instruction param ident - let first_arg = self.get_first_instruction_arg()?; - let params_ident = &first_arg.name; - - // Generate CreateMintsCpi invocation for all mints (no PDA context) - let mints = &self.parsed.light_mint_fields; - let mint_invocation = - LightMintsBuilder::new(mints, params_ident, &self.infra).generate_invocation(); - - // Infrastructure field reference for quote! interpolation - let fee_payer = &self.infra.fee_payer; - - Ok(quote! { - // Build CPI accounts with CPI context enabled (mints use CPI context for batching) - let cpi_accounts = light_sdk::cpi::v2::CpiAccounts::new_with_config( - &self.#fee_payer, - _remaining, - light_sdk::cpi::CpiAccountsConfig::new_with_cpi_context(crate::LIGHT_CPI_SIGNER), - ); - - // Create mints using CreateMintsCpi - #mint_invocation - - Ok(true) - }) - } - - /// Generate LightPreInit body for PDAs only (no mints) - /// After this, compressed addresses are registered - pub fn generate_pre_init_pdas_only(&self) -> Result { + /// Generate PDAs-only body WITHOUT the Ok(true) return. + fn generate_pre_init_pdas_only_body(&self) -> Result { let (compress_blocks, new_addr_idents) = generate_pda_compress_blocks(&self.parsed.rentfree_fields); let rentfree_count = self.parsed.rentfree_fields.len() as u8; - // Get instruction param ident let first_arg = self.get_first_instruction_arg()?; let params_ident = &first_arg.name; - // Infra field references let fee_payer = &self.infra.fee_payer; let compression_config = &self.infra.compression_config; Ok(quote! { - // Build CPI accounts (no CPI context needed for PDAs-only) let cpi_accounts = light_sdk::cpi::v2::CpiAccounts::new( &self.#fee_payer, _remaining, crate::LIGHT_CPI_SIGNER, ); - // Load compression config let compression_config_data = light_sdk::interface::LightConfig::load_checked( &self.#compression_config, &crate::ID )?; - // Collect compressed infos for all rentfree accounts let mut all_compressed_infos = Vec::with_capacity(#rentfree_count as usize); #(#compress_blocks)* - // Execute Light System Program CPI directly with proof use light_sdk::cpi::{InvokeLightSystemProgram, LightCpiInstruction}; light_sdk::cpi::v2::LightSystemProgramCpi::new_cpi( crate::LIGHT_CPI_SIGNER, @@ -348,8 +402,28 @@ impl LightAccountsBuilder { .with_new_addresses(&[#(#new_addr_idents),*]) .with_account_infos(&all_compressed_infos) .invoke(cpi_accounts)?; + }) + } - Ok(true) + /// Generate mints-only body WITHOUT the Ok(true) return. + fn generate_pre_init_mints_only_body(&self) -> Result { + let first_arg = self.get_first_instruction_arg()?; + let params_ident = &first_arg.name; + + let mints = &self.parsed.light_mint_fields; + let mint_invocation = + LightMintsBuilder::new(mints, params_ident, &self.infra).generate_invocation(); + + let fee_payer = &self.infra.fee_payer; + + Ok(quote! { + let cpi_accounts = light_sdk::cpi::v2::CpiAccounts::new_with_config( + &self.#fee_payer, + _remaining, + light_sdk::cpi::CpiAccountsConfig::new_with_cpi_context(crate::LIGHT_CPI_SIGNER), + ); + + #mint_invocation }) } @@ -403,24 +477,4 @@ impl LightAccountsBuilder { } }) } - - /// Check if token accounts or ATAs need finalize code generation. - pub fn needs_token_finalize(&self) -> bool { - TokenAccountsBuilder::new( - &self.parsed.token_account_fields, - &self.parsed.ata_fields, - &self.infra, - ) - .needs_finalize() - } - - /// Generate finalize body for token accounts and ATAs. - pub fn generate_token_finalize_body(&self) -> TokenStream { - TokenAccountsBuilder::new( - &self.parsed.token_account_fields, - &self.parsed.ata_fields, - &self.infra, - ) - .generate_finalize_body() - } } diff --git a/sdk-libs/macros/src/light_pdas/accounts/derive.rs b/sdk-libs/macros/src/light_pdas/accounts/derive.rs index d55e6af909..80c74096cf 100644 --- a/sdk-libs/macros/src/light_pdas/accounts/derive.rs +++ b/sdk-libs/macros/src/light_pdas/accounts/derive.rs @@ -3,18 +3,24 @@ //! This module coordinates code generation by combining: //! - PDA block generation from `pda.rs` //! - Mint action invocation from `mint.rs` +//! - Token account creation from `token.rs` //! - Parsing results from `parse.rs` //! -//! Design for mints: -//! - At mint init, we CREATE + DECOMPRESS atomically -//! - After init, the Mint should always be in decompressed/"hot" state +//! Design: ALL account creation happens in pre_init (before instruction handler) //! -//! Flow for PDAs + mints: -//! 1. Pre-init: ALL compression logic executes here +//! Account types handled: +//! - PDAs (compressed accounts) +//! - Mints (compressed mints - CREATE + DECOMPRESS atomically) +//! - Token accounts (vaults for transfers) +//! - ATAs (associated token accounts) +//! +//! Flow: +//! 1. Pre-init: ALL account creation executes here //! a. Write PDAs to CPI context -//! b. Invoke mint_action with decompress + CPI context -//! c. Mint is now "hot" and usable -//! 2. Instruction body: Can use hot Mint (mintTo, transfers, etc.) +//! b. Create mints with decompress + CPI context +//! c. Create token accounts (vaults) +//! d. Create ATAs +//! 2. Instruction body: All accounts available for use (transfers, minting, etc.) //! 3. Finalize: No-op (all work done in pre_init) use proc_macro2::TokenStream; @@ -33,26 +39,15 @@ pub(super) fn derive_light_accounts(input: &DeriveInput) -> Result { @@ -90,10 +85,10 @@ mod tests { let output = result.unwrap().to_string(); - // Verify finalize generates token account creation + // Verify pre_init generates token account creation assert!( - output.contains("LightFinalize"), - "Should generate LightFinalize impl" + output.contains("LightPreInit"), + "Should generate LightPreInit impl" ); assert!( output.contains("CreateTokenAccountCpi"), @@ -111,7 +106,7 @@ mod tests { #[test] fn test_ata_with_init_generates_create_cpi() { - // ATA with init should generate create_associated_token_account_idempotent in finalize + // ATA with init should generate CreateTokenAtaCpi in pre_init let input: DeriveInput = parse_quote! { #[instruction(params: CreateAtaParams)] pub struct CreateAta<'info> { @@ -133,10 +128,10 @@ mod tests { let output = result.unwrap().to_string(); - // Verify finalize generates ATA creation + // Verify pre_init generates ATA creation assert!( - output.contains("LightFinalize"), - "Should generate LightFinalize impl" + output.contains("LightPreInit"), + "Should generate LightPreInit impl" ); assert!( output.contains("CreateTokenAtaCpi"), @@ -145,7 +140,7 @@ mod tests { } #[test] - fn test_token_mark_only_generates_noop_finalize() { + fn test_token_mark_only_generates_no_creation() { // Token without init should NOT generate creation code (mark-only mode) // Mark-only returns None from parsing, so token_account_fields is empty let input: DeriveInput = parse_quote! { @@ -171,16 +166,20 @@ mod tests { "Mark-only should NOT generate CreateTokenAccountCpi" ); - // Should still have LightFinalize but with Ok(()) + // Should still generate both trait impls + assert!( + output.contains("LightPreInit"), + "Should generate LightPreInit impl" + ); assert!( output.contains("LightFinalize"), - "Should still generate LightFinalize impl" + "Should generate LightFinalize impl (no-op)" ); } #[test] fn test_mixed_token_and_ata_generates_both() { - // Mixed token account + ATA should generate both creation codes + // Mixed token account + ATA should generate both creation codes in pre_init let input: DeriveInput = parse_quote! { #[instruction(params: CreateBothParams)] pub struct CreateBoth<'info> { @@ -206,7 +205,11 @@ mod tests { let output = result.unwrap().to_string(); - // Should have both creation types + // Should have both creation types in pre_init + assert!( + output.contains("LightPreInit"), + "Should generate LightPreInit impl" + ); assert!( output.contains("CreateTokenAccountCpi"), "Should generate CreateTokenAccountCpi for vault" diff --git a/sdk-libs/macros/src/light_pdas/accounts/token.rs b/sdk-libs/macros/src/light_pdas/accounts/token.rs index 9fd2911da0..64fac5e472 100644 --- a/sdk-libs/macros/src/light_pdas/accounts/token.rs +++ b/sdk-libs/macros/src/light_pdas/accounts/token.rs @@ -5,8 +5,8 @@ //! //! ## Code Generation //! -//! Token accounts and ATAs are created in `LightFinalize` (after instruction logic) -//! because they may depend on data set during the instruction handler. +//! Token accounts and ATAs are created in `LightPreInit` (before instruction logic) +//! so they are available for use during the instruction handler (transfers, etc.). //! //! - **Token Accounts**: Use `CreateTokenAccountCpi` with PDA signing //! - **ATAs**: Use `CreateTokenAtaCpi` with `idempotent()` builder @@ -204,15 +204,18 @@ impl<'a> TokenAccountsBuilder<'a> { } /// Check if any token accounts or ATAs need to be created. - pub fn needs_finalize(&self) -> bool { + pub fn needs_creation(&self) -> bool { self.token_account_fields.iter().any(|f| f.has_init) || self.ata_fields.iter().any(|f| f.has_init) } - /// Generate the finalize body for token accounts and ATAs. - pub fn generate_finalize_body(&self) -> TokenStream { - if !self.needs_finalize() { - return quote! { Ok(()) }; + /// Generate token account and ATA creation code for pre_init. + /// + /// Returns None if no token accounts or ATAs need to be created. + /// Otherwise returns the CPI code (without Ok() return). + pub fn generate_pre_init_token_creation(&self) -> Option { + if !self.needs_creation() { + return None; } // Generate token account creation code @@ -229,17 +232,15 @@ impl<'a> TokenAccountsBuilder<'a> { .filter_map(|f| generate_ata_cpi(f, self.infra)) .collect(); - quote! { + Some(quote! { // Get system program from the struct's system_program field let __system_program = self.system_program.to_account_info(); - // Create token accounts + // Create token accounts (in pre_init so they're available for instruction logic) #(#token_account_cpis)* - // Create ATAs + // Create ATAs (in pre_init so they're available for instruction logic) #(#ata_cpis)* - - Ok(()) - } + }) } }