Skip to main content

sol_parser_sdk/core/
common_filler.rs

1use crate::core::events::*;
2use crate::core::invoke_context::{InvokeContext, InvokeLookup};
3use solana_sdk::pubkey::Pubkey;
4use std::collections::HashMap;
5use yellowstone_grpc_proto::prelude::{Transaction, TransactionStatusMeta};
6
7/// `get_fees` / `get_fees_with_quote_mint`: discriminator followed by is_pump_pool.
8/// Other instructions in the fee program do not carry this parameter.
9#[inline]
10pub(crate) fn pumpswap_is_pump_pool_from_fee_instruction(data: &[u8]) -> Option<bool> {
11    const GET_FEES: [u8; 8] = [231, 37, 126, 85, 207, 91, 63, 52];
12    const GET_FEES_WITH_QUOTE: [u8; 8] = [154, 237, 138, 92, 162, 2, 162, 187];
13    let discriminator: [u8; 8] = data.get(..8)?.try_into().ok()?;
14    if discriminator != GET_FEES && discriminator != GET_FEES_WITH_QUOTE {
15        return None;
16    }
17    match data.get(8) {
18        Some(0) => Some(false),
19        Some(1) => Some(true),
20        _ => None,
21    }
22}
23
24#[inline]
25fn set_pumpswap_is_pump_pool_from_fees_ix<L: InvokeLookup + ?Sized>(
26    meta: &TransactionStatusMeta,
27    transaction: &Option<Transaction>,
28    program_invokes: &L,
29    pool: Pubkey,
30    user: Pubkey,
31    buy: bool,
32    is_pump_pool: &mut bool,
33) {
34    // Standalone fee queries carry no pool identity. Only a direct fee CPI
35    // belonging to the uniquely matched swap can enrich this event.
36    let Some(swaps) = program_invokes.get_invokes(&crate::grpc::program_ids::PUMPSWAP_PROGRAM) else {
37        return;
38    };
39    let Some(swap) = crate::core::account_dispatcher::find_pumpswap_trade_invoke(
40        swaps, meta, transaction, pool, user, buy,
41    ) else {
42        return;
43    };
44    let Some(group) = meta.inner_instructions.iter().find(|g| g.index == swap.0 as u32) else {
45        return;
46    };
47    let (start, depth) = if swap.1 < 0 {
48        (0, 1)
49    } else {
50        let Some(depth) = group.instructions.get(swap.1 as usize).and_then(|ix| ix.stack_height) else {
51            return;
52        };
53        (swap.1 as usize + 1, depth)
54    };
55    let Some(fees) = program_invokes.get_invokes(&crate::grpc::program_ids::PUMPSWAP_FEES_PROGRAM) else {
56        return;
57    };
58    let Some(child_depth) = depth.checked_add(1) else {
59        return;
60    };
61    let mut flag = None;
62    for (inner_index, ix) in group.instructions.iter().enumerate().skip(start) {
63        let Some(height) = ix.stack_height else {
64            // Without invocation depth, parentage is unknown.
65            return;
66        };
67        if height <= depth {
68            break;
69        }
70        if height == child_depth {
71            if let Some(value) = pumpswap_is_pump_pool_from_fee_instruction(&ix.data) {
72                if !fees.contains(&(swap.0, inner_index as i32)) {
73                    continue;
74                }
75                if flag.is_some_and(|previous| previous != value) {
76                    return;
77                }
78                flag = Some(value);
79            }
80        }
81    }
82    if let Some(flag) = flag {
83        *is_pump_pool = flag;
84    }
85}
86
87#[inline]
88fn fill_data_with_lookup<L: InvokeLookup + ?Sized>(
89    event: &mut DexEvent,
90    meta: &TransactionStatusMeta,
91    transaction: &Option<Transaction>,
92    program_invokes: &L,
93) {
94    match event {
95        DexEvent::PumpSwapBuy(ref mut e) => {
96            set_pumpswap_is_pump_pool_from_fees_ix(
97                meta,
98                transaction,
99                program_invokes,
100                e.pool,
101                e.user,
102                true,
103                &mut e.is_pump_pool,
104            );
105        }
106        DexEvent::PumpSwapSell(ref mut e) => {
107            set_pumpswap_is_pump_pool_from_fees_ix(
108                meta,
109                transaction,
110                program_invokes,
111                e.pool,
112                e.user,
113                false,
114                &mut e.is_pump_pool,
115            );
116        }
117        _ => {}
118    }
119}
120
121pub fn fill_data(
122    event: &mut DexEvent,
123    meta: &TransactionStatusMeta,
124    transaction: &Option<Transaction>,
125    program_invokes: &HashMap<Pubkey, Vec<(i32, i32)>>,
126) {
127    fill_data_with_lookup(event, meta, transaction, program_invokes);
128}
129
130#[inline]
131pub(crate) fn fill_data_with_invoke_context(
132    event: &mut DexEvent,
133    meta: &TransactionStatusMeta,
134    transaction: &Option<Transaction>,
135    program_invokes: &InvokeContext,
136) {
137    fill_data_with_lookup(event, meta, transaction, program_invokes);
138}
139
140/// Fill PumpFun user post-transaction balances without any RPC calls.
141///
142/// `token_balance` is the raw base-unit balance of the event's `associated_user` token account.
143/// `sol_balance` is the lamport balance of the event's `user` account. A token account closed by
144/// the transaction has a final balance of zero.
145#[inline]
146pub fn fill_token_balances(
147    event: &mut DexEvent,
148    meta: &TransactionStatusMeta,
149    transaction: &Option<Transaction>,
150) {
151    let trade = match event {
152        DexEvent::PumpFunTrade(e)
153        | DexEvent::PumpFunBuy(e)
154        | DexEvent::PumpFunSell(e)
155        | DexEvent::PumpFunBuyExactSolIn(e) => e,
156        _ => return,
157    };
158
159    if let Some(user_index) = account_index(transaction, meta, &trade.user) {
160        trade.sol_balance = meta.post_balances.get(user_index).copied();
161    }
162
163    if let Some(balance) = post_token_balance(transaction, meta, &trade.associated_user) {
164        trade.token_balance = balance;
165    }
166}
167
168#[inline]
169fn account_index(
170    transaction: &Option<Transaction>,
171    meta: &TransactionStatusMeta,
172    account: &Pubkey,
173) -> Option<usize> {
174    if *account == Pubkey::default() {
175        return None;
176    }
177
178    let message = transaction.as_ref()?.message.as_ref()?;
179    message
180        .account_keys
181        .iter()
182        .chain(meta.loaded_writable_addresses.iter())
183        .chain(meta.loaded_readonly_addresses.iter())
184        .position(|key| key.as_slice() == account.as_ref())
185}
186
187#[inline]
188fn post_token_balance(
189    transaction: &Option<Transaction>,
190    meta: &TransactionStatusMeta,
191    account: &Pubkey,
192) -> Option<Option<u64>> {
193    if *account == Pubkey::default() {
194        return None;
195    }
196
197    let message = transaction.as_ref()?.message.as_ref()?;
198    let static_len = message.account_keys.len();
199    let writable_len = meta.loaded_writable_addresses.len();
200
201    let matches_account = |balance: &yellowstone_grpc_proto::prelude::TokenBalance| {
202        let index = balance.account_index as usize;
203        let key = if index < static_len {
204            message.account_keys.get(index)
205        } else if index < static_len + writable_len {
206            meta.loaded_writable_addresses.get(index - static_len)
207        } else {
208            meta.loaded_readonly_addresses.get(index - static_len - writable_len)
209        };
210
211        key.is_some_and(|key| key.as_slice() == account.as_ref())
212    };
213
214    if let Some(balance) = meta.post_token_balances.iter().find(|balance| matches_account(balance))
215    {
216        return Some(crate::grpc::transaction_meta::try_token_balance_raw_amount(balance));
217    }
218
219    meta.pre_token_balances.iter().any(matches_account).then_some(Some(0))
220}
221
222pub fn get_instruction_data<'a>(
223    meta: &'a TransactionStatusMeta,
224    transaction: &'a Option<Transaction>,
225    index: &(i32, i32), // (outer_index, inner_index)
226) -> Option<&'a [u8]> {
227    let data = if index.1 >= 0 {
228        meta.inner_instructions
229            .iter()
230            .find(|i| i.index == index.0 as u32)?
231            .instructions
232            .get(index.1 as usize)?
233            .data
234            .as_slice()
235    } else {
236        transaction.as_ref()?.message.as_ref()?.instructions.get(index.0 as usize)?.data.as_slice()
237    };
238    Some(data)
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use yellowstone_grpc_proto::prelude::{Message, TokenBalance, UiTokenAmount};
245
246    fn token_balance(account_index: u32, amount: u64) -> TokenBalance {
247        TokenBalance {
248            account_index,
249            ui_token_amount: Some(UiTokenAmount {
250                amount: amount.to_string(),
251                ..Default::default()
252            }),
253            ..Default::default()
254        }
255    }
256
257    fn pumpfun_event(user: Pubkey, associated_user: Pubkey) -> DexEvent {
258        DexEvent::PumpFunBuy(PumpFunTradeEvent {
259            user,
260            associated_user,
261            is_buy: true,
262            ..Default::default()
263        })
264    }
265
266    fn transaction(accounts: &[Pubkey]) -> Option<Transaction> {
267        Some(Transaction {
268            message: Some(Message {
269                account_keys: accounts.iter().map(|key| key.to_bytes().to_vec()).collect(),
270                ..Default::default()
271            }),
272            ..Default::default()
273        })
274    }
275
276    #[test]
277    fn fills_pumpfun_user_post_transaction_balances() {
278        let user = Pubkey::new_unique();
279        let associated_user = Pubkey::new_unique();
280        let mut event = pumpfun_event(user, associated_user);
281        let meta = TransactionStatusMeta {
282            pre_balances: vec![50_000, 1_000],
283            post_balances: vec![40_000, 1_000],
284            pre_token_balances: vec![token_balance(1, 10)],
285            post_token_balances: vec![token_balance(1, 35)],
286            ..Default::default()
287        };
288
289        fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
290
291        let DexEvent::PumpFunBuy(trade) = event else { panic!("expected PumpFun buy") };
292        assert_eq!(trade.sol_balance, Some(40_000));
293        assert_eq!(trade.token_balance, Some(35));
294    }
295
296    #[test]
297    fn missing_token_balance_side_means_zero_for_created_or_closed_account() {
298        let user = Pubkey::new_unique();
299        let associated_user = Pubkey::new_unique();
300        let tx = transaction(&[user, associated_user]);
301
302        let mut created = pumpfun_event(user, associated_user);
303        let created_meta = TransactionStatusMeta {
304            pre_balances: vec![50, 0],
305            post_balances: vec![40, 1],
306            post_token_balances: vec![token_balance(1, 25)],
307            ..Default::default()
308        };
309        fill_token_balances(&mut created, &created_meta, &tx);
310        let DexEvent::PumpFunBuy(created) = created else { unreachable!() };
311        assert_eq!(created.token_balance, Some(25));
312
313        let mut closed = pumpfun_event(user, associated_user);
314        let closed_meta = TransactionStatusMeta {
315            pre_balances: vec![40, 1],
316            post_balances: vec![45, 0],
317            pre_token_balances: vec![token_balance(1, 25)],
318            ..Default::default()
319        };
320        fill_token_balances(&mut closed, &closed_meta, &tx);
321        let DexEvent::PumpFunBuy(closed) = closed else { unreachable!() };
322        assert_eq!(closed.token_balance, Some(0));
323    }
324
325    #[test]
326    fn ignores_non_pumpfun_events() {
327        let mut event = DexEvent::Error("unchanged".to_string());
328
329        fill_token_balances(&mut event, &TransactionStatusMeta::default(), &None);
330
331        assert!(matches!(event, DexEvent::Error(ref message) if message == "unchanged"));
332    }
333
334    #[test]
335    fn resolves_user_token_account_from_loaded_writable_addresses() {
336        let user = Pubkey::new_unique();
337        let associated_user = Pubkey::new_unique();
338        let mut event = pumpfun_event(user, associated_user);
339        let tx = transaction(&[user]);
340        let meta = TransactionStatusMeta {
341            pre_balances: vec![50, 1],
342            post_balances: vec![40, 1],
343            loaded_writable_addresses: vec![associated_user.to_bytes().to_vec()],
344            pre_token_balances: vec![token_balance(1, 10)],
345            post_token_balances: vec![token_balance(1, 35)],
346            ..Default::default()
347        };
348
349        fill_token_balances(&mut event, &meta, &tx);
350
351        let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
352        assert_eq!(trade.token_balance, Some(35));
353    }
354
355    #[test]
356    fn ignores_out_of_range_token_balance_account_index() {
357        let user = Pubkey::new_unique();
358        let associated_user = Pubkey::new_unique();
359        let mut event = pumpfun_event(user, associated_user);
360        let meta = TransactionStatusMeta {
361            pre_balances: vec![50, 1],
362            post_balances: vec![40, 1],
363            pre_token_balances: vec![token_balance(99, 10)],
364            post_token_balances: vec![token_balance(99, 35)],
365            ..Default::default()
366        };
367
368        fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
369
370        let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
371        assert_eq!(trade.token_balance, None);
372    }
373
374    #[test]
375    fn malformed_post_token_amount_is_unknown_not_zero() {
376        let user = Pubkey::new_unique();
377        let associated_user = Pubkey::new_unique();
378        let mut event = pumpfun_event(user, associated_user);
379        let meta = TransactionStatusMeta {
380            pre_balances: vec![50, 1],
381            post_balances: vec![40, 1],
382            post_token_balances: vec![TokenBalance {
383                account_index: 1,
384                ui_token_amount: Some(UiTokenAmount {
385                    amount: "invalid".to_string(),
386                    ..Default::default()
387                }),
388                ..Default::default()
389            }],
390            ..Default::default()
391        };
392
393        fill_token_balances(&mut event, &meta, &transaction(&[user, associated_user]));
394
395        let DexEvent::PumpFunBuy(trade) = event else { unreachable!() };
396        assert_eq!(trade.token_balance, None);
397        assert_eq!(trade.sol_balance, Some(40));
398    }
399}
400
401#[cfg(test)]
402mod fee_instruction_layout_tests {
403    use super::pumpswap_is_pump_pool_from_fee_instruction;
404
405    #[test]
406    fn pump_pool_flag_is_independent_of_market_cap_and_other_fee_instructions() {
407        for disc in [[231, 37, 126, 85, 207, 91, 63, 52], [154, 237, 138, 92, 162, 2, 162, 187]] {
408            for (flag, market_cap) in [(0u8, 1u128), (1, 0)] {
409                let mut data = disc.to_vec();
410                data.push(flag);
411                data.extend_from_slice(&market_cap.to_le_bytes());
412                assert_eq!(pumpswap_is_pump_pool_from_fee_instruction(&data), Some(flag != 0));
413            }
414            assert_eq!(pumpswap_is_pump_pool_from_fee_instruction(&disc), None);
415        }
416        let mut unknown = vec![0; 9];
417        unknown[8] = 1;
418        assert_eq!(pumpswap_is_pump_pool_from_fee_instruction(&unknown), None);
419    }
420}
421
422#[cfg(test)]
423mod review_fee_scope_regressions {
424    use super::*;
425    use yellowstone_grpc_proto::prelude::{CompiledInstruction, InnerInstruction, InnerInstructions, Message};
426
427    #[test]
428    fn fees_only_enrich_their_own_swap_and_unknown_parentage_preserves_logs() {
429        let pool0 = Pubkey::new_unique();
430        let pool1 = Pubkey::new_unique();
431        let user = Pubkey::new_unique();
432        let swap = |pool_index| CompiledInstruction {
433            accounts: std::iter::once(pool_index).chain(std::iter::repeat(2).take(22)).collect(),
434            data: crate::instr::pump_amm::discriminators::BUY.to_vec(),
435            ..Default::default()
436        };
437        let fee = |flag| {
438            let mut data = vec![231, 37, 126, 85, 207, 91, 63, 52];
439            data.push(flag);
440            data
441        };
442        let tx = Some(Transaction {
443            message: Some(Message {
444                account_keys: vec![pool0.to_bytes().to_vec(), pool1.to_bytes().to_vec(), user.to_bytes().to_vec()],
445                instructions: vec![swap(0), swap(1), CompiledInstruction { data: fee(1), ..Default::default() }],
446                ..Default::default()
447            }),
448            ..Default::default()
449        });
450        let mut meta = TransactionStatusMeta {
451            inner_instructions: vec![
452                InnerInstructions { index: 0, instructions: vec![InnerInstruction { data: fee(1), stack_height: Some(2), ..Default::default() }] },
453                InnerInstructions { index: 1, instructions: vec![InnerInstruction { data: fee(0), stack_height: Some(2), ..Default::default() }] },
454            ],
455            ..Default::default()
456        };
457        let invokes = HashMap::from([
458            (crate::grpc::program_ids::PUMPSWAP_PROGRAM, vec![(0, -1), (1, -1)]),
459            (crate::grpc::program_ids::PUMPSWAP_FEES_PROGRAM, vec![(0, 0), (1, 0), (2, -1)]),
460        ]);
461        let mut flag = false;
462        set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool0, user, true, &mut flag);
463        assert!(flag);
464        set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool1, user, true, &mut flag);
465        assert!(!flag, "last standalone query must not override the other pool");
466        meta.inner_instructions.clear();
467        set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool0, user, true, &mut flag);
468        assert!(!flag, "standalone query cannot be attributed to this pool");
469        meta.inner_instructions = vec![InnerInstructions { index: 1, instructions: vec![InnerInstruction { data: fee(0), stack_height: None, ..Default::default() }] }];
470        flag = true;
471        set_pumpswap_is_pump_pool_from_fees_ix(&meta, &tx, &invokes, pool1, user, true, &mut flag);
472        assert!(flag, "preserve a log-derived flag without parentage evidence");
473    }
474}