From 47bc81f025a4bc184644a783f630c79fe1339b41 Mon Sep 17 00:00:00 2001 From: ysq Date: Mon, 8 Sep 2025 17:33:53 +0800 Subject: [PATCH] refactor: enhance associated token account handling with custom token program support - Rename create_associated_token_account_fast to create_associated_token_account_idempotent_fast - Upgrade get_associated_token_address_fast to support custom token_program_id parameter - Update ATA cache mechanism to include token_program_id - Adopt new ATA functions in Bonk and PumpFun instruction builders - Improve token program compatibility and cache efficiency --- src/common/fast_fn.rs | 26 +++++++++++++++------- src/instruction/bonk.rs | 24 +++++++++++--------- src/instruction/pumpfun.rs | 45 +++++++++++++++++++++++--------------- 3 files changed, 59 insertions(+), 36 deletions(-) diff --git a/src/common/fast_fn.rs b/src/common/fast_fn.rs index 69ad526..1b1b94c 100644 --- a/src/common/fast_fn.rs +++ b/src/common/fast_fn.rs @@ -6,7 +6,7 @@ use solana_sdk::{ pubkey::Pubkey, }; use spl_associated_token_account::{ - get_associated_token_address, ID as ASSOCIATED_TOKEN_PROGRAM_ID, + get_associated_token_address_with_program_id, ID as ASSOCIATED_TOKEN_PROGRAM_ID, }; use std::num::NonZeroUsize; @@ -61,7 +61,7 @@ where // --------------------- Associated Token Account --------------------- -pub fn create_associated_token_account_fast( +pub fn create_associated_token_account_idempotent_fast( payer: &Pubkey, owner: &Pubkey, mint: &Pubkey, @@ -78,7 +78,8 @@ pub fn create_associated_token_account_fast( // 使用缓存获取指令 get_cached_instruction(cache_key, || { // 使用缓存的方式获取 Associated Token Address - let associated_token_address = get_associated_token_address_fast(owner, mint); + let associated_token_address = + get_associated_token_address_with_program_id_fast(owner, mint, token_program); // 创建 Associated Token Account 指令 // 参考 spl_associated_token_account::instruction::create_associated_token_account 的实现 @@ -92,7 +93,7 @@ pub fn create_associated_token_account_fast( crate::constants::SYSTEM_PROGRAM_META, AccountMeta::new_readonly(*token_program, false), // Token程序（只读，非签名者） ], - data: vec![], // ATA创建指令不需要额外数据 + data: vec![1], } }) } @@ -143,6 +144,7 @@ where struct AtaCacheKey { wallet_address: Pubkey, token_mint_address: Pubkey, + token_program_id: Pubkey, } /// 全局 ATA 缓存，用于存储 Associated Token Address 计算结果 @@ -150,12 +152,16 @@ static ATA_CACHE: Lazy>> = Lazy::new(|| RwLock::new(CLruCache::new(NonZeroUsize::new(MAX_ATA_CACHE_SIZE).unwrap()))); /// 获取缓存的 Associated Token Address，如果不存在则计算并缓存 -pub fn get_associated_token_address_fast( +pub fn get_associated_token_address_with_program_id_fast( wallet_address: &Pubkey, token_mint_address: &Pubkey, + token_program_id: &Pubkey, ) -> Pubkey { - let cache_key = - AtaCacheKey { wallet_address: *wallet_address, token_mint_address: *token_mint_address }; + let cache_key = AtaCacheKey { + wallet_address: *wallet_address, + token_mint_address: *token_mint_address, + token_program_id: *token_program_id, + }; // 尝试从缓存中获取（使用读锁） { @@ -166,7 +172,11 @@ pub fn get_associated_token_address_fast( } // 缓存未命中，计算新的 ATA - let ata = get_associated_token_address(wallet_address, token_mint_address); + let ata = get_associated_token_address_with_program_id( + wallet_address, + token_mint_address, + token_program_id, + ); // 将计算结果存入缓存（使用写锁） { diff --git a/src/instruction/bonk.rs b/src/instruction/bonk.rs index ab9fac7..93c7195 100755 --- a/src/instruction/bonk.rs +++ b/src/instruction/bonk.rs @@ -40,14 +40,18 @@ impl InstructionBuilder for BonkInstructionBuilder { let pool_state = get_pool_pda(¶ms.mint, &crate::constants::WSOL_TOKEN_ACCOUNT).unwrap(); // Create user token accounts - let user_base_token_account = spl_associated_token_account::get_associated_token_address( - ¶ms.payer.pubkey(), - ¶ms.mint, - ); - let user_quote_token_account = spl_associated_token_account::get_associated_token_address( - ¶ms.payer.pubkey(), - &crate::constants::WSOL_TOKEN_ACCOUNT, - ); + let user_base_token_account = + crate::common::fast_fn::get_associated_token_address_with_program_id_fast( + ¶ms.payer.pubkey(), + ¶ms.mint, + &protocol_params.mint_token_program, + ); + let user_quote_token_account = + crate::common::fast_fn::get_associated_token_address_with_program_id_fast( + ¶ms.payer.pubkey(), + &crate::constants::WSOL_TOKEN_ACCOUNT, + &crate::constants::TOKEN_PROGRAM, + ); // Get pool token accounts let base_vault_account = get_vault_pda(&pool_state, ¶ms.mint).unwrap(); @@ -109,7 +113,7 @@ impl InstructionBuilder for BonkInstructionBuilder { // Create buy instruction let accounts = vec![ solana_sdk::instruction::AccountMeta::new(params.payer.pubkey(), true), // Payer (signer) - accounts::AUTHORITY_META, // Authority (readonly) + accounts::AUTHORITY_META, // Authority (readonly) accounts::GLOBAL_CONFIG_META, // Global Config (readonly) solana_sdk::instruction::AccountMeta::new_readonly( protocol_params.platform_config, @@ -243,7 +247,7 @@ impl InstructionBuilder for BonkInstructionBuilder { // Create sell instruction let accounts = vec![ solana_sdk::instruction::AccountMeta::new(params.payer.pubkey(), true), // Payer (signer) - accounts::AUTHORITY_META, // Authority (readonly) + accounts::AUTHORITY_META, // Authority (readonly) accounts::GLOBAL_CONFIG_META, // Global Config (readonly) solana_sdk::instruction::AccountMeta::new_readonly( protocol_params.platform_config, diff --git a/src/instruction/pumpfun.rs b/src/instruction/pumpfun.rs index a92319e..f1bc842 100755 --- a/src/instruction/pumpfun.rs +++ b/src/instruction/pumpfun.rs @@ -57,7 +57,7 @@ impl InstructionBuilder for PumpFunInstructionBuilder { let mut instructions = Vec::with_capacity(2); // Create associated token account - instructions.push(crate::common::fast_fn::create_associated_token_account_fast( + instructions.push(crate::common::fast_fn::create_associated_token_account_idempotent_fast( ¶ms.payer.pubkey(), ¶ms.payer.pubkey(), ¶ms.mint, @@ -75,13 +75,16 @@ impl InstructionBuilder for PumpFunInstructionBuilder { } else { bonding_curve.account }; - let associated_bonding_curve = if protocol_params.associated_bonding_curve - == Pubkey::default() - { - crate::common::fast_fn::get_associated_token_address_fast(&bonding_curve, ¶ms.mint) - } else { - protocol_params.associated_bonding_curve - }; + let associated_bonding_curve = + if protocol_params.associated_bonding_curve == Pubkey::default() { + crate::common::fast_fn::get_associated_token_address_with_program_id_fast( + &bonding_curve, + ¶ms.mint, + &crate::constants::TOKEN_PROGRAM, + ) + } else { + protocol_params.associated_bonding_curve + }; let accounts: [AccountMeta; 16] = [ global_constants::GLOBAL_ACCOUNT_META, @@ -90,9 +93,10 @@ impl InstructionBuilder for PumpFunInstructionBuilder { AccountMeta::new(bonding_curve, false), AccountMeta::new(associated_bonding_curve, false), AccountMeta::new( - crate::common::fast_fn::get_associated_token_address_fast( + crate::common::fast_fn::get_associated_token_address_with_program_id_fast( ¶ms.payer.pubkey(), ¶ms.mint, + &crate::constants::TOKEN_PROGRAM, ), false, ), @@ -139,9 +143,10 @@ impl InstructionBuilder for PumpFunInstructionBuilder { } else { return Err(anyhow!("Amount token is required")); }; - let ata = crate::common::fast_fn::get_associated_token_address_fast( + let ata = crate::common::fast_fn::get_associated_token_address_with_program_id_fast( ¶ms.payer.pubkey(), ¶ms.mint, + &crate::constants::TOKEN_PROGRAM, ); let creator_vault_pda = protocol_params.creator_vault; let creator = get_creator(&creator_vault_pda); @@ -168,13 +173,16 @@ impl InstructionBuilder for PumpFunInstructionBuilder { } else { bonding_curve.account }; - let associated_bonding_curve = if protocol_params.associated_bonding_curve - == Pubkey::default() - { - crate::common::fast_fn::get_associated_token_address_fast(&bonding_curve, ¶ms.mint) - } else { - protocol_params.associated_bonding_curve - }; + let associated_bonding_curve = + if protocol_params.associated_bonding_curve == Pubkey::default() { + crate::common::fast_fn::get_associated_token_address_with_program_id_fast( + &bonding_curve, + ¶ms.mint, + &crate::constants::TOKEN_PROGRAM, + ) + } else { + protocol_params.associated_bonding_curve + }; let accounts: [AccountMeta; 14] = [ global_constants::GLOBAL_ACCOUNT_META, @@ -183,9 +191,10 @@ impl InstructionBuilder for PumpFunInstructionBuilder { AccountMeta::new(bonding_curve, false), AccountMeta::new(associated_bonding_curve, false), AccountMeta::new( - crate::common::fast_fn::get_associated_token_address_fast( + crate::common::fast_fn::get_associated_token_address_with_program_id_fast( ¶ms.payer.pubkey(), ¶ms.mint, + &crate::constants::TOKEN_PROGRAM, ), false, ),