sol_parser_sdk/core/
common_filler.rs

1use crate::core::invoke_context::{InvokeContext, InvokeLookup};
2use crate::{core::events::*, instr::read_bool};
3use solana_sdk::pubkey::Pubkey;
4use std::collections::HashMap;
5use yellowstone_grpc_proto::prelude::{Transaction, TransactionStatusMeta};
6
7#[inline]
8fn set_pumpswap_is_pump_pool_from_fees_ix<L: InvokeLookup + ?Sized>(
9    meta: &TransactionStatusMeta,
10    transaction: &Option<Transaction>,
11    program_invokes: &L,
12    is_pump_pool: &mut bool,
13) {
14    if let Some(invoke) = program_invokes
15        .get_invokes(&crate::grpc::program_ids::PUMPSWAP_FEES_PROGRAM)
16        .and_then(|v| v.last())
17    {
18        if let Some(data) = get_instruction_data(meta, transaction, invoke) {
19            *is_pump_pool = read_bool(data, 9).unwrap_or_default();
20        }
21    }
22}
23
24#[inline]
25fn fill_data_with_lookup<L: InvokeLookup + ?Sized>(
26    event: &mut DexEvent,
27    meta: &TransactionStatusMeta,
28    transaction: &Option<Transaction>,
29    program_invokes: &L,
30) {
31    match event {
32        DexEvent::PumpSwapBuy(ref mut e) => {
33            set_pumpswap_is_pump_pool_from_fees_ix(
34                meta,
35                transaction,
36                program_invokes,
37                &mut e.is_pump_pool,
38            );
39        }
40        DexEvent::PumpSwapSell(ref mut e) => {
41            set_pumpswap_is_pump_pool_from_fees_ix(
42                meta,
43                transaction,
44                program_invokes,
45                &mut e.is_pump_pool,
46            );
47        }
48        _ => {}
49    }
50}
51
52pub fn fill_data(
53    event: &mut DexEvent,
54    meta: &TransactionStatusMeta,
55    transaction: &Option<Transaction>,
56    program_invokes: &HashMap<Pubkey, Vec<(i32, i32)>>,
57) {
58    fill_data_with_lookup(event, meta, transaction, program_invokes);
59}
60
61#[inline]
62pub(crate) fn fill_data_with_invoke_context(
63    event: &mut DexEvent,
64    meta: &TransactionStatusMeta,
65    transaction: &Option<Transaction>,
66    program_invokes: &InvokeContext,
67) {
68    fill_data_with_lookup(event, meta, transaction, program_invokes);
69}
70
71/// Fill PumpFun user post-transaction balances without any RPC calls.
72///
73/// `token_balance` is the raw base-unit balance of the event's `associated_user` token account.
74/// `sol_balance` is the lamport balance of the event's `user` account. A token account closed by
75/// the transaction has a final balance of zero.
76#[inline]
77pub fn fill_token_balances(
78    event: &mut DexEvent,
79    meta: &TransactionStatusMeta,
80    transaction: &Option<Transaction>,
81) {
82    let trade = match event {
83        DexEvent::PumpFunTrade(e)
84        | DexEvent::PumpFunBuy(e)
85        | DexEvent::PumpFunSell(e)
86        | DexEvent::PumpFunBuyExactSolIn(e) => e,
87        _ => return,
88    };
89
90    if let Some(user_index) = account_index(transaction, meta, &trade.user) {
91        trade.sol_balance = meta.post_balances.get(user_index).copied();
92    }
93
94    if let Some(balance) = post_token_balance(transaction, meta, &trade.associated_user) {
95        trade.token_balance = balance;
96    }
97}
98
99#[inline]
100fn account_index(
101    transaction: &Option<Transaction>,
102    meta: &TransactionStatusMeta,
103    account: &Pubkey,
104) -> Option<usize> {
105    if *account == Pubkey::default() {
106        return None;
107    }
108
109    let message = transaction.as_ref()?.message.as_ref()?;
110    message
111        .account_keys
112        .iter()
113        .chain(meta.loaded_writable_addresses.iter())
114        .chain(meta.loaded_readonly_addresses.iter())
115        .position(|key| key.as_slice() == account.as_ref())
116}
117
118#[inline]
119fn post_token_balance(
120    transaction: &Option<Transaction>,
121    meta: &TransactionStatusMeta,
122    account: &Pubkey,
123) -> Option<Option<u64>> {
124    if *account == Pubkey::default() {
125        return None;
126    }
127
128    let message = transaction.as_ref()?.message.as_ref()?;
129    let static_len = message.account_keys.len();
130    let writable_len = meta.loaded_writable_addresses.len();
131
132    let matches_account = |balance: &yellowstone_grpc_proto::prelude::TokenBalance| {
133        let index = balance.account_index as usize;
134        let key = if index < static_len {
135            message.account_keys.get(index)
136        } else if index < static_len + writable_len {
137            meta.loaded_writable_addresses.get(index - static_len)
138        } else {
139            meta.loaded_readonly_addresses.get(index - static_len - writable_len)
140        };
141
142        key.is_some_and(|key| key.as_slice() == account.as_ref())
143    };
144
145    if let Some(balance) = meta.post_token_balances.iter().find(|balance| matches_account(balance))
146    {
147        return Some(crate::grpc::transaction_meta::try_token_balance_raw_amount(balance));
148    }
149
150    meta.pre_token_balances.iter().any(matches_account).then_some(Some(0))
151}
152
153pub fn get_instruction_data<'a>(
154    meta: &'a TransactionStatusMeta,
155    transaction: &'a Option<Transaction>,
156    index: &(i32, i32), // (outer_index, inner_index)
157) -> Option<&'a [u8]> {
158    let data = if index.1 >= 0 {
159        meta.inner_instructions
160            .iter()
161            .find(|i| i.index == index.0 as u32)?
162            .instructions
163            .get(index.1 as usize)?
164            .data
165            .as_slice()
166    } else {
167        transaction.as_ref()?.message.as_ref()?.instructions.get(index.0 as usize)?.data.as_slice()
168    };
169    Some(data)
170}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175    use yellowstone_grpc_proto::prelude::{Message, TokenBalance, UiTokenAmount};
176
177    fn token_balance(account_index: u32, amount: u64) -> TokenBalance {
178        TokenBalance {
179            account_index,
180            ui_token_amount: Some(UiTokenAmount {
181                amount: amount.to_string(),
182                ..Default::default()
183            }),
184            ..Default::default()
185        }
186    }
187
188    fn pumpfun_event(user: Pubkey, associated_user: Pubkey) -> DexEvent {
189        DexEvent::PumpFunBuy(PumpFunTradeEvent {
190            user,
191            associated_user,
192            is_buy: true,
193            ..Default::default()
194        })
195    }
196
197    fn transaction(accounts: &[Pubkey]) -> Option<Transaction> {
198        Some(Transaction {
199            message: Some(Message {
200                account_keys: accounts.iter().map(|key| key.to_bytes().to_vec()).collect(),
201                ..Default::default()
202            }),
203            ..Default::default()
204        })
205    }
206
207    #[test]
208    fn fills_pumpfun_user_post_transaction_balances() {
209        let user = Pubkey::new_unique();
210        let associated_user = Pubkey::new_unique();
211        let mut event = pumpfun_event(user, associated_user);
212        let meta = TransactionStatusMeta {
213            pre_balances: vec![50_000, 1_000],
214            post_balances: vec![40_000, 1_000],
215            pre_token_balances: vec![token_balance(1, 10)],
216            post_token_balances: vec![token_balance(1, 35)],
217            ..Default::default()
218        };
219
220        fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
221
222        let DexEvent::PumpFunBuy(trade) = event else { panic!("expected PumpFun buy") };
223        assert_eq!(trade.sol_balance, Some(40_000));
224        assert_eq!(trade.token_balance, Some(35));
225    }
226
227    #[test]
228    fn missing_token_balance_side_means_zero_for_created_or_closed_account() {
229        let user = Pubkey::new_unique();
230        let associated_user = Pubkey::new_unique();
231        let tx = transaction(&[user, associated_user]);
232
233        let mut created = pumpfun_event(user, associated_user);
234        let created_meta = TransactionStatusMeta {
235            pre_balances: vec![50, 0],
236            post_balances: vec![40, 1],
237            post_token_balances: vec![token_balance(1, 25)],
238            ..Default::default()
239        };
240        fill_token_balances(&mut created, &created_meta, &tx);
241        let DexEvent::PumpFunBuy(created) = created else { unreachable!() };
242        assert_eq!(created.token_balance, Some(25));
243
244        let mut closed = pumpfun_event(user, associated_user);
245        let closed_meta = TransactionStatusMeta {
246            pre_balances: vec![40, 1],
247            post_balances: vec![45, 0],
248            pre_token_balances: vec![token_balance(1, 25)],
249            ..Default::default()
250        };
251        fill_token_balances(&mut closed, &closed_meta, &tx);
252        let DexEvent::PumpFunBuy(closed) = closed else { unreachable!() };
253        assert_eq!(closed.token_balance, Some(0));
254    }
255
256    #[test]
257    fn ignores_non_pumpfun_events() {
258        let mut event = DexEvent::Error("unchanged".to_string());
259
260        fill_token_balances(&mut event, &TransactionStatusMeta::default(), &None);
261
262        assert!(matches!(event, DexEvent::Error(ref message) if message == "unchanged"));
263    }
264
265    #[test]
266    fn resolves_user_token_account_from_loaded_writable_addresses() {
267        let user = Pubkey::new_unique();
268        let associated_user = Pubkey::new_unique();
269        let mut event = pumpfun_event(user, associated_user);
270        let tx = transaction(&[user]);
271        let meta = TransactionStatusMeta {
272            pre_balances: vec![50, 1],
273            post_balances: vec![40, 1],
274            loaded_writable_addresses: vec![associated_user.to_bytes().to_vec()],
275            pre_token_balances: vec![token_balance(1, 10)],
276            post_token_balances: vec![token_balance(1, 35)],
277            ..Default::default()
278        };
279
280        fill_token_balances(&mut event, &meta, &tx);
281
282        let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
283        assert_eq!(trade.token_balance, Some(35));
284    }
285
286    #[test]
287    fn ignores_out_of_range_token_balance_account_index() {
288        let user = Pubkey::new_unique();
289        let associated_user = Pubkey::new_unique();
290        let mut event = pumpfun_event(user, associated_user);
291        let meta = TransactionStatusMeta {
292            pre_balances: vec![50, 1],
293            post_balances: vec![40, 1],
294            pre_token_balances: vec![token_balance(99, 10)],
295            post_token_balances: vec![token_balance(99, 35)],
296            ..Default::default()
297        };
298
299        fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
300
301        let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
302        assert_eq!(trade.token_balance, None);
303    }
304
305    #[test]
306    fn malformed_post_token_amount_is_unknown_not_zero() {
307        let user = Pubkey::new_unique();
308        let associated_user = Pubkey::new_unique();
309        let mut event = pumpfun_event(user, associated_user);
310        let meta = TransactionStatusMeta {
311            pre_balances: vec![50, 1],
312            post_balances: vec![40, 1],
313            post_token_balances: vec![TokenBalance {
314                account_index: 1,
315                ui_token_amount: Some(UiTokenAmount {
316                    amount: "invalid".to_string(),
317                    ..Default::default()
318                }),
319                ..Default::default()
320            }],
321            ..Default::default()
322        };
323
324        fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
325
326        let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
327        assert_eq!(trade.token_balance, None);
328        assert_eq!(trade.sol_balance, Some(40));
329    }
330}
sol_parser_sdk/core/common_filler.rs

sol_parser_sdk/core/
common_filler.rs
