sol_parser_sdk/core/
common_filler.rs

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

sol_parser_sdk/core/
common_filler.rs
