gavindiaz/sol-trade-sdk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: enhance associated token account handling with custom token program support

Browse Source 
- 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
This commit is contained in:
ysq
2025-09-08 17:33:53 +08:00
parent 93c133a9ea
commit 47bc81f025
3 changed files with 59 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/fast_fn.rs
+18 -8
View File
@@ -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<RwLock<CLruCache<AtaCacheKey, Pubkey>>> =
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,
);
// 将计算结果存入缓存（使用写锁）
{
src/instruction/bonk.rs
+14 -10
View File
@@ -40,14 +40,18 @@ impl InstructionBuilder for BonkInstructionBuilder {
let pool_state = get_pool_pda(&params.mint, &crate::constants::WSOL_TOKEN_ACCOUNT).unwrap();
// Create user token accounts
let user_base_token_account = spl_associated_token_account::get_associated_token_address(
&params.payer.pubkey(),
&params.mint,
);
let user_quote_token_account = spl_associated_token_account::get_associated_token_address(
&params.payer.pubkey(),
&crate::constants::WSOL_TOKEN_ACCOUNT,
);
let user_base_token_account =
crate::common::fast_fn::get_associated_token_address_with_program_id_fast(
&params.payer.pubkey(),
&params.mint,
&protocol_params.mint_token_program,
);
let user_quote_token_account =
crate::common::fast_fn::get_associated_token_address_with_program_id_fast(
&params.payer.pubkey(),
&crate::constants::WSOL_TOKEN_ACCOUNT,
&crate::constants::TOKEN_PROGRAM,
);
// Get pool token accounts
let base_vault_account = get_vault_pda(&pool_state, &params.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,
src/instruction/pumpfun.rs
+27 -18
View File
@@ -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(
&params.payer.pubkey(),
&params.payer.pubkey(),
&params.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, &params.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,
&params.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(
&params.payer.pubkey(),
&params.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(
&params.payer.pubkey(),
&params.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, &params.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,
&params.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(
&params.payer.pubkey(),
&params.mint,
&crate::constants::TOKEN_PROGRAM,
),
false,
),